LMCache Token Dropping 实战指南：用 SDK 编辑 KV Cache 提升 vLLM 解码吞吐-创锋一号

LMCache Token Dropping 实战指南：用 SDK 编辑 KV Cache 提升 vLLM 解码吞吐

【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache

导读

长提示词会产生巨大的 KV Cache，占用 GPU 显存并压缩 batch 容量，进而拖累解码吞吐。本文基于 LMCache 仓库中的 examples/token_dropping 示例，完整介绍Token Dropping技术：通过 LMCache SDK 将请求的 KV Cache 取回、按算法丢弃冗余 token、再写回供 vLLM 解码，从而在不改 vLLM 推理内核的前提下把解码吞吐提升 1.5–1.7x。读完本文，你将掌握 Random Dropping、SnapKV、R-KV 三种策略及三种 R-KV 加速变体的原理与取舍、从源码安装 LMCache 的完整步骤、为 vLLM 打补丁以导出 query 中间张量的方法，以及 SDKbatch/request/context的核心 API 调用链。

背景：为什么需要 Token Dropping

长提示词（如长文档问答）生成的 KV Cache 会吃掉大量 GPU 显存，直接限制 batch 中能容纳的请求数量。batch 变小意味着解码阶段（decode）的吞吐下降。要提升解码吞吐，就需要让更多的请求同时挤进一个 batch——而压缩每个请求 KV Cache 的体量是最直接的途径。

Token dropping（Token 丢弃）正是这样一种技术：按某种重要性度量，从每个请求的 KV Cache 中挑选并丢弃一部分 token（本示例中默认丢弃一半），从而缩小每个请求的 KV Cache、提高 decode 吞吐。文档实测的吞吐收益为1.5–1.7x，并且示例同时验证：一个好的 Token Dropping 算法（示例选用SnapKV）不仅不会损伤生成准确率，甚至可能小幅提升准确率——因为丢弃的往往是冗余信息，反而抑制了注意力噪声。

这些示例全部通过LMCache SDK以离线（offline）方式完成：SDK 负责retrieve（取回）请求的缓存张量、调用用户提供的丢弃函数modify（修改）、再store（写回）供 vLLM 解码。用户只需要编写 token dropping 函数本身，SDK 的 batch / stream API 负责其余编排。

示例总览：三种策略与三种 R-KV 加速变体

examples/token_dropping/目录下共有 6 个 notebook，按“是否依赖 query 张量”分为两类：

Notebook	策略	需要 query 张量？
random_token_dropping.ipynb	随机丢弃一部分历史 token，只使用 KV Cache	否
snapkv_token_dropping.ipynb	SnapKV：保留开头窗口、结尾窗口，以及被 recent-window query 最关注的那些 token	是
rkv_token_dropping.ipynb	R-KV：在 SnapKV 重要性项的基础上减去一个redundancy（冗余）项，token 只有“既被关注、又提供新信息”才被保留。三个优化变体见 rkv_variants/	是

R-KV 的冗余项需要比较所有 key 两两之间的相似度，时间与内存均为O(n²)。rkv_token_dropping.ipynb 按论文原始方式实现；rkv_variants/ 下的三个独立 notebook 则从不同角度把它做便宜：

Notebook	冗余项实现	精确？	VRAM 开销
rkv_variant_cpu_exact.ipynb	重写为一次内积，n×n 矩阵始终不会物化	是	无
rkv_variant_gpu_exact.ipynb	同样的改写，按行分块（tile）在 GPU 上执行，四个版本中最快	是	约 1.55 GiB
rkv_variant_buffered_cpu.ipynb	只在 chunk 内部比较 key，矩阵变为块对角	否，约 96% token 一致	无

此外 snapkv_colab.ipynb 与 rkv_colab.ipynb 是面向 Google Colab GPU T4 的演示版，使用更小的模型与更短的评测数据集以适配 T4 显存（每个 notebook 大约需要 15 分钟运行）。

验证思路：如果只想快速复现“decode 吞吐提升”，从最简单的 random_token_dropping.ipynb 入手即可——它不需要 query 张量，也就不需要vLLM 补丁和transfer_query开关。

前置条件

GPU：要观察到 token dropping 放大 decode batch 的效果，需要结合请求数量一起调节--gpu-memory-utilization。示例在单张 RTX 6000 PRO 上运行。
LMCache：SDK 传输层可选用共享内存（shared memory）或 pickle。本示例使用共享内存；若要传输 query 张量，启动 LMCache 时必须传入--enable transfer_query（对应 lmcache/v1/multiprocess/config.py 中的实验性传输模块列表）。
打过补丁的 vLLM：详见下文“vLLM 补丁”。

从源码安装 LMCache

示例依赖最近dev分支的新特性（--enable transfer_query、engine-driven transfer），因此必须从源码安装，而不是安装已发布的 wheel：

uv venv --python 3.12 && source .venv/bin/activate uv pip install torch # 先装 torch uv pip install -e . --no-build-isolation # 再构建 LMCache

几个常见的坑（文档与源码共同印证）：

必须先装 torch，且必须传--no-build-isolation。csrc/下的原生 CUDA 扩展在编译时要链接已安装 torch 的头文件；不加--no-build-isolation时 pip 会在隔离环境里构建，看不到 torch，编译直接失败。
nvcc（CUDA toolkit）必须与 torch 的 CUDA 版本一致，否则.cu文件编译报错。CUDA 12 与 13 的依赖会分别从requirements/（如 requirements/cuda12_core.txt、requirements/cuda13_core.txt）自动选择。
每次拉取dev分支后要重新构建。-e（editable）安装让 Python 改动立即生效，但原生扩展不会——如果上游改了csrc/或原生枚举/内核，需要重跑uv pip install -e . --no-build-isolation，否则可能遇到AttributeError: ... 'EngineKVFormat' has no attribute ...（典型的.so过期症状，EngineKVFormat定义于 csrc/engine_kv_format.h）。
用uv pip（或 venv 内的pip），不要用系统pip——系统 pip 可能被标记为 externally-managed（PEP 668）而拒绝安装。

构建变体速查：NO_NATIVE_EXT=1（纯 Python，无扩展）、NO_GPU_EXT=1 ... --no-build-isolation（仅 CPU C++）、BUILD_WITH_HIP=1（AMD ROCm/HIP）。

vLLM 补丁：导出中间张量（query）

许多 Token Dropping 算法（如 SnapKV）需要 query 张量来给 token 的重要性打分，而 vLLM 默认不向 KV connector 暴露中间张量。仓库提供了一份约 10 行的补丁 vllm-export-intermediate-tensors.diff 解决这个问题，它只改动两个 Python 文件、无需重新编译，打完补丁重启 vLLM 即可。

从 diff 内容看，改动分两处：

在vllm/distributed/kv_transfer/kv_connector/v1/base.py的KVConnectorBase_V1上新增一个默认返回False的transfer_intermediate_tensors属性；
在vllm/model_executor/layers/attention/kv_transfer_utils.py的maybe_transfer_kv_layer中，当该属性为True时，把save_kv_layer调用的绑定参数（bound_args.arguments）作为intermediate_tensors一并传下去。

安装前请确认 vLLM 版本在补丁验证过的范围（0.23.0 到 0.25.1）。

有 vLLM 源码 checkout（git 树）：直接用git apply：

cd /path/to/vllm git apply /path/to/LMCache/examples/token_dropping/vllm-export-intermediate-tensors.diff

从 wheel 安装的 vLLM（pip install vllm，没有 git 树）：直接把 diff 应用到已安装的包：

VLLM_DIR=$(python -c "import vllm, os; print(os.path.dirname(vllm.__file__))") patch -p1 -d "$VLLM_DIR/.." < /path/to/LMCache/examples/token_dropping/vllm-export-intermediate-tensors.diff

启动 vLLM 时，通过在 connector 配置中加入lmcache.mp.transfer_intermediate_tensors激活该路径。默认情况下，QRingBuffer（用于暂存 query 张量的临时 staging buffer，实现在 lmcache/sdk/qringbuffer.py）的容量可以容纳2 次 forward pass的 query 张量，可通过"lmcache.mp.q.ring_depth":2调整：

--kv-transfer-config '{ "kv_connector": "LMCacheMPConnector", "kv_role": "kv_both", "kv_connector_extra_config": { "lmcache.mp.port": 6555, "lmcache.mp.transfer_intermediate_tensors": true, "lmcache.mp.q.ring_depth":2 } }'

需要强调：random dropping 示例不需要这份补丁和这个开关——它只依赖 KV Cache，是演示 decode 吞吐提升的最简路径。注意transfer_query（LMCache 侧--enable开关）与lmcache.mp.transfer_intermediate_tensors（vLLM connector 侧配置）是同一能力的两个端点：前者在 LMCache server 侧启用 query 传输模块，后者在 vLLM 侧开启中间张量导出。

数据集

Notebook 默认加载 Hugging Face 数据集raniayu/token-dropping-demo：从 LongBench-v2 中挑选出30 个prompt 长度最接近 10240 token 的样本（按 Qwen3-8B tokenizer 计算）。

Google Colab 版本（适配 T4 显存）使用更短的数据集raniayu/token-dropping-demo-short：从 ehovy/race 中挑选10 个prompt 长度最接近 1024 token、且前缀上下文各不相同的样本。

SDK 工作原理：retrieve → modify → store

示例用到的 SDK 能力全部围绕“取回—修改—写回”三段式展开，对应源码为：

retrieve：LMCacheRequestStream.retrieve() 以当前 token 序列为 key，按 chunk 对齐取回缓存的 KV 张量（形状[2, L, T, D]），并轮询等待缓存就绪；
modify：LMCacheRequestStream.modify_kv() 取出每个缓存类型对应的张量后，调用用户提供的fn(tensors, tokens)——也就是ModifyFnType（定义在 lmcache/sdk/context.py，签名(Mapping[CacheKind, Tensor], Sequence[int]) -> (Tensor, Sequence[int])），返回修改后的(new_kv, new_tokens)；chunk 对齐之外的尾部 token 会被暂存进_suffix_tokens，在下一次generate()时拼接回去；
store：LMCacheRequestStream.update() 将编辑后的 KV 通过context.store()写回（写入时校验len(tokens) == kv.shape[2]），并重置流的 token 序列与 done 标记，让 vLLM 从新的 KV 开始解码。

Batch 编排由 LMCacheBatchedStream 提供：prefill()对每个流做一次max_tokens=1的预填充并输出输入吞吐/TTFT 指标；modify()把 KV 编辑函数并行应用到所有流并报告耗时；decode()并行解码并输出输出吞吐/TPOT 指标。三个步骤的度量分别通过get_perf_metrics()（mode 为"prefill"或"decode"）汇总成终端 Metrics 报告。

另外，LMCacheSDKCacheKind（lmcache/sdk/cache_kind.py）区分两种可缓存张量族：KV与QUERY，其中 QUERY 类型在 server 侧以{model_name}##query作为命名空间前缀。SnapKV/R-KV 这类需要 query 的算法，会在创建上下文时同时连接 KV 与 QUERY 两种 kind（通过 lmcache/sdk/init.py 的connect()分发到kvcache.connect/qcache.connect）。

一个值得注意的底层细节在 examples/token_dropping/utils.py 的rerotate_k_cache：丢弃 token 后，剩余 key 在 cache 中的位置改变了，必须把它们的 RoPE 从旧位置“重旋转”到新位置，否则位置编码与 token 实际位置错位。该函数按层处理[2, L, T, D]张量，先分离旧 RoPE 再做新 RoPE，并且为了数值精度：旋转算术在fp32下执行、最后只 round 一次回 cache 的 dtype——注释指出若直接在 bfloat16（8 位尾数）中做来回旋转，误差可高达 2.5e-1（对标准差为 3 的 key），而 fp32 路径把往返误差压到约 1e-6。utils.py还提供了面向 vLLM 的流式post_completion（需要 vLLM 以--return-tokens-as-token-ids启动）、LongBench-v2 的答案抽取与打分（extract_answer_choice/score_answers），用于量化丢弃前后准确率变化。

最小复现路径

按上文“从源码安装 LMCache”创建 venv 并安装；
启动 LMCache server，带--enable transfer_query（若跑 SnapKV/R-KV）；
按上文给 vLLM 打补丁（random dropping 可跳过），以--kv-transfer-config传入 connector 配置（含lmcache.mp.transfer_intermediate_tensors: true、lmcache.mp.q.ring_depth）启动 vLLM，并确认 LMCache 的 KV Cache 与 query 张量已注册；
打开 random_token_dropping.ipynb（或 SnapKV/R-KV 版本），加载token-dropping-demo数据集，运行 prefill → modify（传入丢弃函数）→ decode 三步，对比吞吐与准确率指标；
调节--gpu-memory-utilization与请求数量，观察 batch 扩容与解码吞吐提升。

结语

Token Dropping 的完整链路——SDK 取回缓存、用户函数修改、写回 vLLM 解码——展示了 LMCache SDK “离线编辑 KV Cache” 的通用能力：它不侵入 vLLM 推理内核，用户只需提供一段纯 PyTorch 的修改函数。仓库示例同时给出了算法上的深度（SnapKV 与 R-KV、三种复杂度/显存各异的 R-KV 变体）与工程上的完整度（从源码安装、vLLM 补丁、RoPE 重旋转、精度处理、指标评测）。在长上下文推理场景下，这套方法为“显存不够但想塞下更大 batch”提供了开箱即用的实践路径。

【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

企业官网建设流程全解析