资讯动态

vllm-ascend KVPP(KV Layer Parallelism)设计详解:按层拆分持久化 KV 缓存与整层预取

发布时间:2026/10/6 1:54:38 来源:尧图企业网站定制
人工智能大模型模型推理服务AscendCANN【免费下载链接】vllm-ascendCommunity maintained hardware plugin for vLLM on Huawei Ascend项目地址https://gitcode.com/gh_mirrors/vl/vllm-ascend点击查看免费下载KVPPKV Layer Parallelism是 vllm-ascend 为 MLA/SFA 类模型设计的一种 KV 缓存并行化方案它不再让每个 TP rank 复制完整的持久化历史 KV而是按层layer把缓存分片到 TP 组内的不同 rank 上计算到某层时再通过整层广播取回。本文围绕 kvpp.md 的设计文档展开结合仓库内 ascend_config.py、kv_cache_placement.py、kvpp.py 等源码讲清它的背景动机、内存预算公式、广播预取机制、支持边界以及如何通过一行additional_config开启。读完本文你将掌握KVPP 解决了什么问题、为什么用整层广播 双 scratch buffer 预取一层就能成倍扩大长上下文容量、每 rank 的物理内存预算如何计算、哪些模型与并行组合可用、哪些配置会直接报错以及完整的启动命令示例。1. 背景为什么 MLA/SFA 的持久 KV 缓存会被复制在张量并行Tensor ParallelTP下MLAMulti-head Latent Attention与 SFASparse Flash Attention类模型的历史 KV 缓存会在 TP rank 之间复制一份replicate。以 TP2 为例同一段历史上下文在每个 rank 上都保留一份完整缓存直接带来两方面的代价长上下文容量受限KV 缓存按 TP 倍数重复占用显存可容纳的 token 数被压缩并发请求受限缓存副本越多留给并发请求 batch 的空间越小。KVPP 的核心思路与常规 KV 分片不同它**按层by layer在 TP 组内分配持久缓存的所有权而不是按 head 或按序列分片。某个 rank 只持久保存属于自己的那些层的主 KV 缓存当模型执行到某一层、需要该层历史缓存参与注意力计算时拥有该层的 rank 把整层缓存广播broadcast**给组内其他 rank计算完成后再由 owner 保留更新后的缓存供下一轮 forward 继续使用。从实现上看这一策略在 KVPPConfig.from_vllm_config 中被触发只有additional_config[enable_kvpp]为true时KVPP 的组规模才被计算为tensor_parallel_size * prefill_context_parallel_size在不启用 DCP 时MLA 缓存在 PCP 的 KV gather 之后同样被复制因此层归属要覆盖整个复制域。值得强调的是KVPP只改变 KV 缓存的物理存放位置与传输方式不改变模型的执行拓扑。模型执行仍然保留原有的 TP/EP/PP 配置调度器scheduler仍然持有完整、逻辑的缓存规格cache specification与 block 编号。也就是说上层逻辑缓存管理完全不变变的只是底层物理存储的布局和广播时机——这保证了与 vLLM 原有缓存管理、前缀缓存prefix caching、chunked prefill 等机制的兼容。2. 总体设计层归属 双 scratch buffer 预取一层KVPP 的设计由三个可叠加的机制组成层归属Layer Ownership每个 PP stage 把本地的目标层target layer按层号排序后切分成连续、均衡的分区分配给 TP 组内的 owner rank两个可复用的 scratch buffer每个 rank 只为自己拥有的层分配持久存储处理其他层时复用两个 scratch 缓冲预取一层prefetch one layer ahead计算当前层的同时用独立的传输流预先广播下一层。设计文档给出的拓扑示意如下TP group within one PP stage Owner rank: persistent main KV indexer KV for its layers | -- Full-layer broadcast -- Other ranks: scratch buffer | -- Wait -- Compute current layer Prefetch next layer从源码可以验证双 scratch的取值并非随意在 kv_cache_placement.py 中明确定义了# One buffer for the current layer and one for the next layers prefetch. KVPP_SCRATCH_BUFFER_COUNT 2目标层执行序号execution ordinal交替选择两个 scratch buffer保证当前层正在被计算与下一层正在被预取写入可以分别使用独立的存储互不覆盖。3. 层归属Layer Ownership如何给层找主人3.1 分区规则每个 Pipeline ParallelPPstage 先收集本 stage 负责的本地目标层按层号排序然后连续、均衡地分给 KVPP 组内的 owner rank。一层的主 KV、indexer KV 以及量化quantization组件被捆绑为一个整体bundle只有一个 owner不允许把同一层的不同组件拆给不同 rank。对应源码是 map_kvpp_layers_to_owners层按extract_layer_index分桶并排序用divmod(len(layer_indices), kvpp_size)得到base, remainder前remainder个 rank 多分 1 层即partition_size base int(owner_rank remainder)每个 rank 拿到的是一段连续的层区间同 index 下的所有 cache 名主 KV、indexer 等一起归属同一个 owner。之所以要连续分区而非轮询打散是为了让广播的传输模式更规整同一 rank 连续拥有相邻层且各 rank 的层归属在两个 rank 上保持一致同时便于每 rank 的持久分配按层形成连续内存段。3.2 MTP 缓存单独处理MTPMulti-Token Prediction等 draft 层的缓存不属于目标层分区它们仍在每个需要它们的 rank 上独立分配因此 draft 层不会出现在layer_owner_ranks映射中。这一点在 map_kvpp_layers_to_owners 里通过find_draft_layers(...)显式排除并在 register_kvpp_draft_layers 中从 loader 处读取真实的 draft 层名而非靠层号推断同时校验 draft 层不得与目标层共享 KV 缓存KVPP does not support draft layers sharing target KV caches.。3.3 每 rank 的存储角色owner rank为该层分配持久存储persistent main KV indexer KVforward 结束后保留更新后的缓存其他 rank把收到的整层广播写入对应 scratch buffer仅用于本次计算计算结束后该 scratch 可被下一层复用。在 allocate_kvpp_cache 中可以看到plan.layer_owner_ranks.get(name)为None或等于当前 rank 时分配独立持久 buffer否则按target_index % KVPP_SCRATCH_BUFFER_COUNT交替复用两个 scratch bufferscratch 的尺寸取所有目标层 bundle 中的最大值scratch_size max(...)并按 2 MiB 对齐KVPP_BUFFER_ALIGNMENT 2 * 1024 * 1024见 kvpp_cache.py。4. 物理布局与内存预算4.1 物理布局bundle 内部的各组件主 KV、indexer 数据、scale 等在物理上连续存放每个组件包含该层的全部逻辑 block。KVPP 不要求所有层连成一个全局大块——只有一个广播 bundle必须连续这正是广播能直接走一块连续内存的前提。布局严格跟随实际缓存规格cache specification覆盖主 KV、indexer 数据、scale/量化分量以及 LI-C8、SFA-C8 等稀疏缓存布局。对应实现是 build_kvpp_layer_layout按 bundle 内 cache 名顺序、按tensor_sizes中每个分量的每 block 字节数 × num_blocks推进游标cursor得到每个分量的(offset, size)区间build_kvpp_buffer_sizes 则根据AscendSFAIndexerCacheSpec、AscendMLAAttentionSpec、量化 split factor 等逐分量计算每 block 大小。4.2 每 rank 的物理内存预算设计文档给出了核心公式Physical bytes per block Sum of owned target-bundle bytes per block Local MTP-cache bytes per block 2 * Largest target-bundle bytes per block Available blocks floor(Available KV memory / Physical bytes per block)逐项解读owned target-bundle bytes per block本 rank 拥有的所有目标层 bundle 每 block 字节数之和——这是唯一的持久开销Local MTP-cache bytes per block本地 MTP/draft 缓存每个需要的 rank 都要放一份2 × Largest target-bundle bytes per block两个 scratch buffer各按最大目标层 bundle尺寸预留。最终Available blocks取整后即为 KVPP 物理分配所能容纳的 block 数。worker 再把该 block 容量换算成逻辑内存预算交给 planner逻辑缓存管理调度、block 编号、前缀缓存完全不变物理存储则按 KVPP 布局分配。源码中的对应实现是 KVPPPhysicalCachePlan.get_num_blocksfor name, bundle in self.layer_bundles.items(): _, size build_kvpp_layer_layout(bundle, self.tensor_sizes, num_blocks1) owner self.layer_owner_ranks.get(name) if owner is None or owner self.kvpp_rank: persistent_bytes size if owner is not None: scratch_bytes max(scratch_bytes, size) bytes_per_block persistent_bytes KVPP_SCRATCH_BUFFER_COUNT * scratch_bytes return available_bytes // bytes_per_block if bytes_per_block else 0与文档公式一一对应persistent_bytes累加owned bundle 本地 MTPowner 为 None 的 draft 层scratch_bytes取目标层中的最大值并乘以 2。从该公式可以直观看出收益来源设 TP 组规模为 N、层数为 L原本每个 rank 要存放全部 L 层的完整缓存副本KVPP 之后每 rank 只需持久保存约 L/N 层的 bundle其余层用两块 scratch 轮转承载。层数 L 越大、TP 规模 N 越大、单层缓存越大节省的显存越可观。worker 侧对预算的实际换算在 worker.py 的_apply_kvpp_memory_budget中完成。5. 整层广播与预取Full-Layer Broadcast and Prefetch5.1 执行时序带历史缓存的真实请求forward pass 中包含此前已计算过的 token的执行顺序如下forward 开始前立即预取第一层第一个 attention 层的缓存每个 attention 层在 projection 之后、第一次访问/写入本层 KV 缓存之前等待本层广播完成等待完成后立刻发起下一层的预取使下一层广播与当前层计算重叠。无历史缓存的 forward首次 prefill、dummy runwarmup以及 profiling 路径不做历史缓存广播——它们没有可广播的历史数据这是设计文档明确给出的优化。5.2 广播的语义每次广播只在当前 PP stage 的 KVPP 组内进行不会跨 PP stage 或跨 DP 副本owner 将完整 bundle一次性广播包含所有组件、所有已分配的逻辑 block其他 rank 把接收内容写入对应的 scratch buffer计算结束后owner 保留更新后的缓存用于后续 forward。值得注意的一个开销特征广播载荷不做请求级或 active-block 级过滤——只要 block 已分配就整层整块传输广播流量随分配的 block 总数增长而不是随实际命中的 token 数增长。这在长上下文 高并发时会成为主要的通信开销需要在容量收益与通信开销之间权衡。5.3 传输流与事件同步源码级机制预取使用独立的传输流transfer stream通过事件event建立严格的数据依赖与完成保证。核心实现在 BroadcastKVPPTransport.prefetchdef prefetch(self, layer_name, cache_ready, transfer_stream): with torch.npu.stream(transfer_stream): transfer_stream.wait_event(cache_ready) # 依赖等当前流就绪 work dist.broadcast( self._layer_buffers[layer_name], srcself._owner_global_ranks[layer_name], groupself._device_group, async_opTrue, ) work.wait() # 等待广播调用完成 done torch.npu.Event() done.record(transfer_stream) # 在传输流上记录完成事件 # A Future must cover device completion, including the owners source # reads, before attention can overwrite the persistent cache. done.synchronize() # 设备级完成后再返回要点cache_ready事件在当前计算流上记录传输流wait_event(cache_ready)保证广播不会早于缓存数据就绪dist.broadcast(..., async_opTrue)在独立线程/流上发起work.wait()只保证调用层面完成done.record(transfer_stream)之后done.synchronize()保证设备传输真正结束包括 owner 侧源数据读取完毕才从 wait 返回——避免 attention 在广播尚未落盘时就去覆写持久缓存。5.4 预取调度器预取由 KVPPScheduler 驱动schedule_forward(has_history)has_history为真时立即start_layer_prefetch(第一个 attention 层)预取任务提交到单线程ThreadPoolExecutorkvpp-prefetch线程在线程内torch.npu.set_device后走传输流wait_for_layer(layer_name)当前层广播完成后_next_attention_layer_index 1若还有下一层则立即发起下一层预取——这就是计算当前层与预取下一层重叠的直接实现complete_forward()清空_has_history与层游标等待下一轮 forward。执行侧接入在 KVPPRuntimecreate_from_kv_cache依据KVPPConfig.from_vllm_config判断config.size 1才构建运行时通过static_forward_context[name].impl.layerwise_kv_cache_hook scheduler把逐层钩子挂到各 attention 层Model Runner V1 中则在 model_runner_v1.py 用num_computed_tokens 0判定has_history并调用prepare_forwardforward 结束调用complete_forward。5.5 KVPP 通信组如何建立KVPP 组在 parallel_state.py 中建立所有 rank 按[ExternalDP, DP, PP, PCP, TP]张量化后flatten(-2)把 PCP 与 TP 维度合并再按kvpp_size切分出组每个 DP 副本、每个 PP stage 各有一个独立的 KVPP 组组的规模恒等于pcp_size * tp_size。这样层归属的复制域与 MLA 缓存实际被复制的域一致DCP 关闭时PCP 的 prefill gather 也会复制 MLA KV。6. 支持范围与约束Support and Constraints设计文档给出了完整的支持矩阵维度范围模型与执行非混合non-hybridMLA/SFA 模型eager 模式Model Runner V1 与 V2支持的组合TP、EP、PP、chunked prefill、prefix caching前缀缓存、异步调度async scheduling、固定步长 MTPfixed-step MTP缓存布局按实际缓存规格分配包括 LI-C8 与 SFA-C8不支持Graph 执行CUDA Graph 类、PCP、DCP、prefill/decode 分离disaggregated、变步长 MTPvariable-step MTP源码层面的校验逻辑位于 KVPPConfig.validate会在配置阶段直接抛错值得逐一对照DCP 互斥decode_context_parallel_size ! 1时报KVPP and DCP cannot be enabled at the same time.PD 分离限制KVPP 的 PD 分离仅允许MooncakeConnectorV2连接器且 decode-only 节点kv_consumer上必须关闭 KVPPKVPP must be disabled on the decode-only node.执行模式仅支持 eager或 CUDA Graph 模式为PIECEWISEKVPP supports eager execution or PIECEWISE only.模型类型仅支持非混合 MLA 模型KVPP currently supports only non-hybrid MLA models.投机解码仅支持methodmtp或methoddspark且投机 token 数必须固定KVPP currently supports only a fixed number of speculative tokens.DSpark 自适应校验adaptive verification与动态投机长度均被拒绝。这些校验意味着只要配置不满足上述条件服务会启动失败并给出明确原因而不是静默降级便于用户在部署前快速发现不兼容组合。通信换容量的本质KVPP 的本质是用通信换持久缓存容量每个 rank 仍然需要两块按最大目标层尺寸预留的 scratch buffer外加独立的 MTP 缓存持久缓存则从全部层缩减为本地拥有的 L/N 层。因此内存收益取决于层数、TP 规模、每层缓存尺寸三者的组合通信代价取决于分配的 block 总数广播载荷不按请求/active block 过滤。层数越多、TP 规模越大每 rank 的持久份额越小收益越明显而 block 数越大每层广播的载荷越重对互连带宽的要求越高。7. 启用方式UsageKVPP 通过模型启动参数开启开关为additional_config中的enable_kvpp默认false。设计文档给出的最小示例vllm serve model-path \ --tensor-parallel-size 2 \ --enforce-eager \ --additional-config {enable_kvpp: true}要点说明--enforce-eager必须显式给出如第 6 节所述Graph 执行除 PIECEWISE 模式外不被支持不满足时配置校验会直接报错--tensor-parallel-size 2KVPP 组规模默认跟随 TP 规模更精确地说源码中为tensor_parallel_size * prefill_context_parallel_size即 MLA 缓存被复制的复制域不启用 PCP 时即为 TP 规模。组内各 rank 各自持有部分层的持久缓存其余层通过广播取用开启 PP 时每个 PP stage 在自己的 TP 组内独立分配 owner 并广播跨 stage 之间互不干扰开关解析enable_kvpp通过 validate_additional_config_bool 按 pydantic 的 bool 规则解析字符串false/true也会被正确转换随后由 KVPPConfig.from_vllm_config 读取。开启后日志侧可在ascend_config的kvpp_config.size字段确认生效size 1表示 KVPP 已激活见 ascend_config.py 中kvpp_config: KVPPConfig的注入。若需与投机解码MTP/DSpark联用请确保投机 token 数为固定值并注意 MTP 缓存会在各所需 rank 独立分配、不参与层分区。8. 小结KVPP 为 vllm-ascend 上的 MLA/SFA 长上下文场景提供了一个清晰的显存-通信权衡杠杆架构上保持逻辑缓存管理与模型执行拓扑不变仅把持久 KV 按层分散到 TP 组各 rank配以整层广播 双 scratch buffer 预取一层的运行时机制内存上每 rank 的持久开销从全部层降为本地层 两块最大层 scratch 本地 MTP公式与源码get_num_blocks一一对应约束上eager 或 PIECEWISE、非混合 MLA、固定步长投机是硬性前提PCP/DCP/PD 分离/变步长 MTP 会被配置校验直接拒绝使用上一条--additional-config {enable_kvpp: true}即可开启配合--enforce-eager使用。如果你正在为 DeepSeek 等 MLA 架构的长上下文 高并发部署头疼缓存容量KVPP 是把复制 N 份的历史缓存改为每 rank 持有 L/N 层 整层广播的务实选项——具体收益建议结合你的层数、TP 规模与单层缓存尺寸按第 4 节的公式先行估算再评估互连带宽能否承受随 block 数线性增长的广播流量。赞分享人工智能大模型模型推理服务AscendCANN【免费下载链接】vllm-ascendCommunity maintained hardware plugin for vLLM on Huawei Ascend项目地址https://gitcode.com/gh_mirrors/vl/vllm-ascend点击查看免费下载相关推荐跨层 KV 共享Cross-Layer KV Sharing实战从 LLMs-from-scratch 仓库理解 KV 缓存的内存优化跨层 KV 共享Cross Layer KV Sharing实战从 LLMs from scratch 仓库理解 KV 缓存的内存优化 本篇以 LLMs示例工程大模型人工智能vLLM-Ascend 分层与稀疏 KV Cache 卸载Prefill/Decode 分离部署的设计与实践vLLM Ascend 分层与稀疏 KV Cache 卸载Prefill/Decode 分离部署的设计与实践 本文围绕 vLLM Ascend 的分层Lay人工智能大模型模型推理服务AscendCANNSGLang 分层 KV 缓存接入 AIBrix KVCache 实战指南以 AIBrixKVCache 作为 L3 KV 缓存SGLang 分层 KV 缓存接入 AIBrix KVCache 实战指南以 AIBrixKVCache 作为 L3 KV 缓存 本文以仓库文档 python模型推理服务推理引擎人工智能大模型本地部署多模态上一篇MediaCrawler cookie管理维持登录状态的高级技巧下一篇突破线性笔记局限SiYuan图谱视图让知识连接可视化创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑