ARTICLE DETAIL

资讯详情

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

LMCache 中的 KV Cache 压缩:CacheGen 编解码原理、配置方式与源码解析

LMCache 中的 KV Cache 压缩:CacheGen 编解码原理、配置方式与源码解析 LMCache 中的 KV Cache 压缩CacheGen 编解码原理、配置方式与源码解析【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCacheKV Cache 在长上下文、多轮对话等场景下会迅速占据大量显存与远端存储直接制约缓存命中率与加载速度。LMCache 在docs/source/kv_cache_optimizations/compression/index.rst中正式引入了基于 KV Cache 分布特性的压缩算法 CacheGen将 KV Cache 编码为更紧凑的比特流并以可忽略的解码开销换取存储/内存占用与加载速度的大幅改善。读完本文你将掌握 CacheGen 在 LMCache 中的启用方式离线与在线两种配置、底层量化—CDF 估计—算术编码完整链路以及如何在仓库中验证其压缩效果。为什么需要 KV Cache 压缩KV Cache 压缩能够显著减小缓存体积这对以下两方面都有直接收益存储 / 内存占用压缩后的 KV Cache 可以更少地占用本地 CPU 内存、本地磁盘或远端缓存服务Redis、Mooncake 等的空间让同样的容量容纳更多 token 的缓存加载速度更小的字节流意味着在远端缓存读取cache hit 场景时更少的网络传输时间与反序列化开销从而缩短首 Token 延迟。在 LMCache 当前的压缩功能矩阵中官方支持的压缩算法为 CacheGen对应docs/source/kv_cache_optimizations/compression/index.rst中的列表CacheGenCacheGen: KV Cache Compression and Streaming for Fast Large Language Model ServingSIGMETRICS 论文docs/source/kv_cache_optimizations/compression/cachegen.rst对 CacheGen 做了如下定位CacheGen 利用 KV Cache 的分布特性distributional properties将其编码为更紧凑的比特流表示bitstream且解码开销可以忽略不计negligible decoding overhead。需要注意上述文档描述的是 LMCache 的 in-process 模式已标记为 deprecated行为。LMCache 更推荐使用 LMCache MP 模式 以获得更完善的功能支持与性能MP 模式下的 CacheGen 支持尚未实现见 MP 模式 CacheGen 文档其内容为 not implemented yet (coming soon)原始 in-process 实现被保留在 Legacy 部分即本文所对应的 compression/index.rst。启用 CacheGen离线与在线两种配置CacheGen 的配置与 LMCache 的 naive KV cache sharing 几乎完全一致只需做少量改动即可启用。也就是说它作用于远端缓存读写这一条链路上对 KV Cache 进行压缩后写入远端、读取时解压还原。离线推理场景环境变量方式在离线推理offline inference中通过设置环境变量LMCACHE_REMOTE_SERDE来启用# Enable cachgen compression in LMCache os.environ[LMCACHE_REMOTE_SERDE] cachegen在线推理场景配置文件方式在在线推理online inference例如与 vLLM 集成的部署中需要在 LMCache 的配置 YAML 中设置remote_serde字段# Enable cachgen compression in LMCache remote_serde: cachegen配置项在源码中的落点remote_serde是LMCacheEngineConfig的一个标准配置项定义于 lmcache/v1/config.py类型Optional[str]默认值naive即默认使用朴素序列化不启用压缩环境变量转换器str从 lmcache/v1/storage_backend/naive_serde/init.py 的CreateSerde工厂函数可以看到serde_type支持naive、kivi、cachegen三种取值其中elif serde_type cachegen: s, d ( CacheGenSerializer(config, metadata), CacheGenDeserializer(config, metadata), )即cachegen会创建CacheGenSerializer与CacheGenDeserializer这对编解码器传入其他非法值会抛出ValueError。另外两点值得注意在 lmcache/v1/cache_engine.py 与 lmcache/v1/cache_engine.py 中创建远端序列化器的分支要求method必须为cachegen进一步印证了 CacheGen 是当前 in-process 模式下唯一受支持的远端压缩算法其余压缩能力走 MP 模式的 legacy 代码路径。远端后端在创建时会断言config.remote_serde is not None见 lmcache/v1/storage_backend/remote_backend.py也就是说使用远端缓存如 Redis、Mooncake时remote_serde是必经之路而默认值naive保证未显式配置时也能正常工作。前置条件先完成跨实例 KV Cache 共享由于 CacheGen 是远端 KV Cache 共享链路中的序列化器正式启用前需要先按照 share_kv_cache 文档 配置好跨实例 KV Cache 共享。LMCache 支持两类共享方式centralized_sharing通过集中式缓存服务器共享 KV Cachep2p_sharing通过点对点缓存传输共享 KV Cache。CacheGen 的配置与这两种共享方式的配置非常相似只需在既有共享配置之上追加remote_serde: cachegen或设置环境变量即可无需调整其他参数。CacheGen 底层原理源码级编解码链路CacheGen 的核心思想是先对 KV Cache 做按层分组的低位宽量化再基于量化值的统计分布估计累积分布函数CDF最后用 CDF 驱动的算术编码把量化值压成紧凑比特流。编码、解码两条路径分别实现在lmcache/storage_backend/serde/cachegen_encoder.pyCacheGenSerializer.to_bytes()lmcache/storage_backend/serde/cachegen_decoder.pyCacheGenDeserializer.from_bytes()1. 按层分组的量化配置QuantizationSpec量化规格由QuantizationSpec描述cachegen_basics.py包含三个字段字段含义start_layer该量化区间的起始层号含end_layer该量化区间的结束层号不含bins量化 bin 数量bin 数越大量化越精细CacheGenConfig聚合了nlayers总层数、kspecsK 的量化区间列表、vspecsV 的量化区间列表并通过from_model_name()按模型族自动推导默认配置7B 系mistralai/Mistral-7B-Instruct-v0.2、lmsys/longchat-7b-16k、Qwen/Qwen-7B32 层K 前 10 层用 32 bins、其余 16 binsV 前 2 层用 32 bins、其余 16 bins8B 系meta-llama/Llama-3.1-8B-Instruct同为 32 层量化分段与 7B 一致9B 系THUDM/glm-4-9b-chat40 层K 前 10 层 32 bins、其余 16 binsV 前 2 层 32 bins、其余 16 bins源码注释标明该配置 needs tuning for better quality其他模型通过AutoConfig.from_pretrained(model_name)读取num_hidden_layers动态生成层数 10 时 K/V 全部用 32 bins否则按前 10 层 / 前 2 层用 32 bins、其余用 16 bins的规则生成无法识别时抛出ValueError。从配置语义可以看出CacheGen 的量化策略偏好是浅层靠近输入的 K/V 分布更重要或更复杂给予更多 bin32深层则用更少的 bin16从而在保真度与压缩率之间取得平衡。2. 量化Quantization编码器入口to_bytes()接收形状为[num_layers, 2, num_tokens, num_heads, head_size]的 blob KV 张量先通过_split_kv()拆成 K 与 Vreshape 后按第 1 维 unbind得到[num_layers, num_tokens, num_channels]形式。核心量化函数torch_quant_vectorized()cachegen_encoder.py逐层按各自 bin 数做对称量化MAX (bins // 2 - 1)[:, None, None] # [nlayers, 1, 1] max1 torch.amax(torch.abs(input_groups), dim-1, keepdimTrue) factor MAX / max1 xq torch.round(input_groups * factor MAX).to(torch.int8)即对每个通道维度取绝对值的最大值作为缩放因子max1将张量缩放到[-MAX, MAX]区间并 round 成int8。max1作为反量化所需的缩放系数随编码结果一并保存。3. CDF 估计与归一化量化完成后编码器基于量化值的直方图统计每个 (层, 通道) 的符号分布并用torch.cumsum计算累积分布函数CDF。为了让浮点 CDF 能安全喂给整数算术编码器_convert_to_int_and_normalize()cachegen_encoder.py做了关键处理将 CDF 从[0, 1)放大到2^1616 位精度量级若需要归一化先乘2^16 - (Lp - 1)再叠加一个arange(Lp)斜坡保证 CDF严格单调递增——算术编码器要求 CDF 单调否则相同符号会出现重复区间导致解码失败。4. 算术编码与 GPU Kernel真正把量化符号压成比特流的是 GPU 端的算术编码 kernel。编码器按CACHEGEN_GPU_MAX_TOKENS_PER_CHUNK 256cachegen_basics.py将 token 分批处理调用device_ops.calculate_cdf(...)GPU 上直接计算 CDF对应 csrc/cuda/cal_cdf.cudevice_ops.encode_fast_new(...)GPU 算术编码对应 csrc/cuda/ac_enc.cu。每批编码结果以CacheGenGPUBytestreambytestreambytestream_lengthsntokens组织最终封装为CacheGenGPUEncoderOutput并通过 pickle 序列化成bytes写入远端。5. 解码与反量化解码路径CacheGenDeserializer.from_bytes()是编码的逆过程反序列化得到CacheGenGPUEncoderOutput调用device_ops.decode_fast_prefsum(...)对应 csrc/cuda/ac_dec.cu按 CDF 与长度前缀和还原每个 (层, 通道) 的量化符号decode_function_gpu()将解码结果 reshape 回[2, nlayers, ntokens, nchannels]并拆出 K/Vdo_dequantize()cachegen_decoder.py用保存的max_tensors_key/max_tensors_value与各层 bin 数做逆缩放t (t - C) / C * maxtensors最后重新 stack 成[nlayers, 2, ntokens, num_heads, head_size]并转回原始 dtype。解码器内部还复用了get_output_buffer()的预分配缓冲来降低反复分配的抖动并对key_bins/value_bins做设备同步源码注释 #83 记录了多 worker 设备隐式绑定问题的临时修复。6. 与 vLLM 集成时的注意点在 vLLM 集成路径中lmcache/integration/vllm/utils.py 存在如下约束if use_mla and (config.remote_serde ! naive and config.remote_serde is not None):即当模型使用 MLAMulti-head Latent Attention且remote_serde被设置为非naive时会有额外的检查逻辑。这意味着在启用 CacheGen 前需要确认目标模型架构尤其 MLA 类模型与 CacheGen 的兼容性约束。如何验证压缩效果仓库自带的端到端 Benchmark仓库在 tests/benchmarks/test_cachegen.py 提供了 CacheGen 编码/解码的端到端基准测试可以直观验证压缩率与耗时。运行方式pytest tests/benchmarks/test_cachegen.py --benchmark-only该测试的关键设计使用_generate_kv()构造 32 层、8 头、head_size 128 的 bfloat16 合成 KV blob编解码器通过LMCacheEngineConfig.from_defaults(chunk_size...)与LMCacheMetadata(model_namemistralai/Mistral-7B-Instruct-v0.2, ...)构建即走 7B 模型的默认量化配置编码测试对chunk_size分别取[64, 128, 256, 768]四档测量to_bytes()耗时并打印raw / compressed的压缩率解码测试测量from_bytes()耗时测试要求当前 torch 运行时可用torch_dev.is_available()否则整模块 skip——因为 CacheGen 依赖手写的 GPU kernelCUDA/XPU 等。此外tests/test_serde.py 等单测覆盖了序列化器的一致性校验可用于确认 CacheGen 编码-解码往返round-trip的正确性。使用建议与限制小结适用路径CacheGen 作用于远端 KV Cache 读写的序列化链路启用前请先完成 KV Cache 跨实例共享 的配置模式限制本文档对应的 in-process 模式已被标记为 deprecatedMP 模式 的 CacheGen 尚未实现生产环境建议关注 MP 模式路线并留意后续版本硬件依赖CacheGen 的 CDF 计算与算术编解码均在 GPU kernel 中完成csrc/cuda 下的ac_enc.cu、ac_dec.cu、cal_cdf.cu启用前需确认推理设备具备对应运行时模型适配默认量化配置针对常见 7B/8B/9B 模型族做了手工调优其他模型会基于num_hidden_layers自动推导如需更优的精度/压缩权衡可关注后续对量化规格可配置化的演进架构约束MLA 类模型与remote_serde非默认值存在联动检查使用前应核对 vllm 集成代码 中的约束条件。整体来看CacheGen 是 LMCache 在KV Cache 压缩方向上的核心实现其分布感知的按层量化 算术编码思路在压缩率与解码开销之间做了精心取舍通过remote_serde: cachegen一行配置即可接入配合仓库自带的 benchmark 可以快速量化其在目标硬件上的实际收益。【免费下载链接】LMCacheLMCache: Supercharge Your LLM with the Fastest KV Cache Layer项目地址: https://gitcode.com/GitHub_Trending/lm/LMCache创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表