ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

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

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 droppingToken 丢弃正是这样一种技术按某种重要性度量从每个请求的 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.ipynbSnapKV保留开头窗口、结尾窗口以及被 recent-window query 最关注的那些 token是rkv_token_dropping.ipynbR-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 GiBrkv_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 上运行。LMCacheSDK 传输层可选用共享内存shared memory或 pickle。本示例使用共享内存若要传输 query 张量启动 LMCache 时必须传入--enable transfer_query对应 lmcache/v1/multiprocess/config.py 中的实验性传输模块列表。打过补丁的 vLLM详见下文“vLLM 补丁”。从源码安装 LMCache示例依赖最近dev分支的新特性--enable transfer_query、engine-driven transfer因此必须从源码安装而不是安装已发布的 wheeluv 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编译直接失败。nvccCUDA toolkit必须与 torch 的 CUDA 版本一致否则.cu文件编译报错。CUDA 12 与 13 的依赖会分别从requirements/如 requirements/cuda12_core.txt、requirements/cuda13_core.txt自动选择。每次拉取dev分支后要重新构建。-eeditable安装让 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-managedPEP 668而拒绝安装。构建变体速查NO_NATIVE_EXT1纯 Python无扩展、NO_GPU_EXT1 ... --no-build-isolation仅 CPU C、BUILD_WITH_HIP1AMD 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 源码 checkoutgit 树直接用git applycd /path/to/vllm git apply /path/to/LMCache/examples/token_dropping/vllm-export-intermediate-tensors.diff从 wheel 安装的 vLLMpip 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_queryLMCache 侧--enable开关与lmcache.mp.transfer_intermediate_tensorsvLLM 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 能力全部围绕“取回—修改—写回”三段式展开对应源码为retrieveLMCacheRequestStream.retrieve() 以当前 token 序列为 key按 chunk 对齐取回缓存的 KV 张量形状[2, L, T, D]并轮询等待缓存就绪modifyLMCacheRequestStream.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()时拼接回去storeLMCacheRequestStream.update() 将编辑后的 KV 通过context.store()写回写入时校验len(tokens) kv.shape[2]并重置流的 token 序列与 done 标记让 vLLM 从新的 KV 开始解码。Batch 编排由 LMCacheBatchedStream 提供prefill()对每个流做一次max_tokens1的预填充并输出输入吞吐/TTFT 指标modify()把 KV 编辑函数并行应用到所有流并报告耗时decode()并行解码并输出输出吞吐/TPOT 指标。三个步骤的度量分别通过get_perf_metrics()mode 为prefill或decode汇总成终端 Metrics 报告。另外LMCacheSDKCacheKindlmcache/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——注释指出若直接在 bfloat168 位尾数中做来回旋转误差可高达 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),仅供参考
返回列表