KV Cache 管理设计文档【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer1. 概述本文档描述了 LLM 推理框架中 KV Cache 管理的设计与实现。该模块采用Paged Attention技术将 KV Cache 以块Block为单位进行管理并兼容不同 attention 类型的块分配策略实现高效的内存复用和并发请求支持。1.1 核心目标内存效率通过分块管理避免为每个请求预分配完整序列长度的内存并发支持多请求共享有限的 NPU 内存资源灵活适配支持多种注意力机制Full Attention、Sliding Window Attention1.2 模块目录结构executor/core/kv_cache/ ├── __init__.py # 模块导出 ├── kv_cache_manager.py # 顶层协调器 KVCacheManager ├── single_type_kv_cache_manager.py # 单类型管理器 ├── block_pool.py # 底层块池 BlockPool ├── cache_info.py # Cache 元数据结构 └── cache_utils.py # 辅助工具函数2. 三层结构设计2.1 总体分层KV cache 管理采用三层结构KVCacheManager请求级总协调器负责跨类型的一致性分配。SingleTypeKVCacheManager单一 attention type 的逻辑块管理器负责决定某个请求在该类型下需要多少 block、何时回收旧 block、如何组织 block table。BlockPool物理 block 池负责具体 block id 的发放和回收。可以用下图表示2.2 第一层KVCacheManagerKVCacheManager是请求级入口面向调度层提供统一接口。定义位置kv_cache_manager.py核心职责如下统一接收调度层的 slot 申请请求。将一次请求的分配广播到所有single_type_managers。在真正分配前完成一次全局可分配性检查避免部分 attention type 分配成功、部分失败。提供 request 维度的 block 查询与整体释放能力。其核心接口包括allocate_slots(request_id, computed_tokens, num_new_tokens, lookahead_tokens0)get_block_ids(request_id)free(request_id)format_usage() - str返回逐 attention type的used/total汇总字符串例如full:12/100,sliding:8/100由Scheduler._log_step调用用于在每步状态日志中暴露各 type 的 KV 压力。get_contiguous_buf_infos() - (data_ptrs, data_lens, item_lens, attn_types)把 paged-attention 管理的所有物理 cache tensor 扁平化导出供 PD KV 传输使用——按layer_idx升序、再按层内cache原始定义序遍历仅导出needs_blockTrue的 cache。返回四个等长 listdata_ptrs每个 tensor 的data_ptr()、data_lens每个 tensor 总字节数、item_lens每块字节数 block_size × num_head × dim × element_size、attn_types每项的 attention type用于在对端按 type 查 block id。Prefill 与 Decode 各自独立调用返回顺序由模型配置决定两端必然对齐。其中allocate_slots的关键点是”两阶段处理”具体实现可见块分配核心流程预检查阶段。 对每个 Manager 先执行旧块回收再计算本次需要的新块数量若任一 Manager 空闲块不足则整体返回失败。分配阶段。 只有当所有 Manager 都满足条件时才逐个执行真实分配避免部分分配导致资源不一致。2.3 第二层SingleTypeKVCacheManagerSingleTypeKVCacheManager抽象”某一种 attention type 的 block 管理策略”。定义位置single_type_kv_cache_manager.py其内部成员包括attn_type当前 manager 所属的 attention 类型。block_pool该类型对应的物理 block 池。block_size每个 block 容纳的 token 数。max_model_len模型总长度上限。_null_block空洞占位 block用于逻辑对齐。之所以需要单独设置null_block是因为相关算子要求 block table 中必须存在合法的 block 索引因此实现上会额外申请一个 block 作为占位但该 block 不承载真实的 cache 数据读写。req_to_blocksrequest_id - block_id 列表的映射字典记录每个请求使用的block ids。通用接口包括get_num_blocks_to_allocate计算某个请求当前需要新增多少 block。allocate_new_blocks从BlockPool获取真实 block并写入req_to_blocks。remove_skipped_blocks回收已不再需要保留的旧 block识别请求已经滑出有效注意力范围的 block将这些位置替换为_null_block并把对应物理 block 归还给BlockPool。free请求结束时释放该 request 在当前 attention type 下的全部 block。需要子类实现的抽象方法get_num_skipped_tokens(num_computed_tokens)计算可跳过分配的 token 数。基类抛出NotImplementedError子类必须实现此方法FullAttentionManager始终返回 0全注意力机制无需跳过任何 token。SlidingWindowManager返回max(0, num_computed_tokens - sliding_window 1)滑窗外的历史 token 可被跳过。validate_and_build_kwargs(group_entries)【静态方法】 验证并构建 manager 实例化所需的类型特定参数。基类抛出NotImplementedError子类必须实现此方法FullAttentionManager无需额外参数返回空字典{}。SlidingWindowManager从CacheEntry中提取并校验sliding_window参数。根据管理注册表可以选择Attention类对应的 cache 管理策略ATTN_TYPE_MANAGER_MAP { FullAttention: FullAttentionManager, SlidingWindow: SlidingWindowManager, }目前已实现两类具体的派生类型类名说明定义位置FullAttentionManager全注意力块随序列增长线性追加single_type_kv_cache_manager.pySlidingWindowManager滑动窗口注意力支持过期块回收single_type_kv_cache_manager.py2.3.1 FullAttentionManagerFullAttentionManager直接复用基类逻辑特点是逻辑 block 数量等于ceil(num_tokens / block_size)。请求随着 token 增长持续追加物理内存块 block。需实现的抽象方法get_num_skipped_tokens(_num_computed_tokens)返回 0全注意力机制无需跳过任何 token。validate_and_build_kwargs(_group_entries)返回空字典{}无额外参数。2.3.2 SlidingWindowManagerSlidingWindowManager在基类上增加滑窗逻辑目标是在逻辑位置保持对齐的同时只为有效窗口保留真实物理 block。其关键设计点如下由于算子约束逻辑块表block_table长度要求与ceil(kv_len / block_size)对齐因此即使窗口左侧的历史 token 已经失效这部分逻辑 block 位置仍然需要保留所以让失效的历史 block 不再占用真实物理块而是用_null_block占位。prefill 时只为有效窗口部分分配真实 block但会在块表头部补齐若干_null_block。decode 时随着num_computed_tokens增长通过remove_skipped_blocks逐步把滑窗之外的 block 置为_null_block并回收到池中。SlidingWindowManager额外实现了_get_logical_block_layout(num_tokens)get_num_blocks_to_allocate(...)allocate_new_blocks(...)get_num_skipped_tokens(num_computed_tokens)返回max(0, num_computed_tokens - sliding_window 1)。validate_and_build_kwargs(group_entries)【静态方法】提取并校验sliding_window参数。这使得滑窗类 attention 可以兼顾“逻辑位置连续”与“物理存储压缩”。2.4 第三层BlockPoolBlockPool是最底层的物理资源管理器只关心 block id 的生命周期。定义位置block_pool.py内部成员包括num_blocks总 block 数。block_sizeblock 大小。_free_block_queue空闲 block 队列负责按顺序发放 block。_free_block_set空闲 block 集合负责快速校验重复释放。_null_block保留的占位 block。其关键设计点如下使用dequeset双结构队列保证 FIFO 分配顺序集合支持 O(1) 重复释放检测。_null_block该 block 是为了满足算子约束而额外申请的保底占位 block不参与正常 cache 存储可以用于表示逻辑上不存在物理块的位置如滑动窗口的无效区域。其核心行为如下初始化时先创建0 ~ num_blocks-1的 block 队列。从队列头取出一个 block 作为_null_block该 block 不参与正常分配也不承载实际 KV 数据它仅用于满足算子约束下的占位需求。get_new_blocks(n)从空闲队列中依次弹出n个 block 并返回。free_blocks(blocks)将释放的 block 重新追加回空闲队列末尾。3. 块分配核心流程allocate_slotsallocate_slots是整个 KV cache 管理链路里的核心入口它的设计重点是逐请求地进行不同 attention type 的块分配。调度层只需要关心当前请求是否还能继续推进而KVCacheManager负责把这一请求广播到所有SingleTypeKVCacheManager分别计算所需块数、回收可跳过块并在真正分配前完成一次整体可分配性检查。3.1 输入参数含义KVCacheManager.allocate_slots(request_id, computed_tokens, num_new_tokens, lookahead_tokens0)中各入参的含义为request_id请求唯一标识。computed_tokens当前请求已经计算完成的 token 数。num_new_tokens本轮需要继续为该请求预留 slot 的新增 token 数。lookahead_tokens在num_new_tokens之外额外预留的 token 数用于 speculative decoding 等需要提前占位的场景。函数内部首先构造两个关键量total_computed_tokens min(computed_tokens, max_model_len)表示当前已经实际计算过并纳入 Cache 管理的token数可用于判断 skipped block 的 token 数。num_tokens_need_slot min(computed_tokens num_new_tokens lookahead_tokens, max_model_len)表示本次分配后需要 block 覆盖到的总 token 范围。3.3 两阶段实现流程3.3.1 预检查阶段预检查阶段是整个分配逻辑最关键的部分它确保分配前先收敛状态。对KVCacheManager中每个SingleTypeKVCacheManager依次执行remove_skipped_blocks(request_id, total_computed_tokens)如果该 attention type 存在可以释放的无用内存块则先回收这部分 block。get_num_blocks_to_allocate(request_id, num_tokens_need_slot)根据num_tokens_need_slot和当前请求已持有的 block 数计算本轮还需新增多少 block。通用逻辑可以表示为required_block_num ceil(num_tokens_need_slot / block_size) existing_block_num len(req_to_blocks.get(request_id, [])) block_num_to_allocate max(0, required_block_num - existing_block_num)get_num_free_blocks()查看BlookPool中空闲块的数量。若任何一个 manager 的空闲块数少于需要新增的块数则顶层立刻返回False该请求在本次迭代中不会发生任何新的 block 分配。3.3.2 分配阶段只有全部 manager 都通过预检查才进入分配阶段此时KVCacheManager会再次遍历 managers从BlockPool申请块并更新 request 对应的 block table。allocate_new_blocks(request_id, block_num_to_allocate, num_tokens_need_slot)3.4 allocate_slots 时序图通过时序图展示这两个阶段的协作关系4. 辅助模块辅助模块包括cache_info.py和cache_utils.py。cache_info.py定义 KV cache 的元数据结构cache_utils.py负责初始化校验、资源分配以及运行期输入准备。4.1 CacheInfocache_info.py定义了三层元数据结构定义位置cache_info.pyCacheEntry描述单个 cache 条目包括cache_name、attn_type、dim、num_head、dtype、tensor_setter和needs_blockSliding Window 场景下还包含sliding_window。LayerCacheInfo按 layer 聚合多个CacheEntry。ModelCacheInfo汇总全模型 cache 信息包括num_layers、block_size和所有层的LayerCacheInfo。CacheEntry.attn_type决定 cache 条目映射到哪一类SingleTypeKVCacheManager。分组维度是attn_type因此本框架支持单层同时包含多个 cache 条目也支持单层同时包含多种 cache 类型。ModelCacheInfo在模型侧构建。执行链路为ModelWorker._get_cache_info()调用模型对象的get_cache_info()由具体模型实现返回自身的 cache 元数据。以 DeepSeek-R1 的 modeling_deepseek.py 为例get_cache_info()会遍历self.model.layers从每层self_attn中取出nope_cache和rope_cache并构造LayerCacheInfo。每层包含两个CacheEntrynope_cacheattn_typeFullAttentiondimself.config.kv_lora_ranknum_head1。rope_cacheattn_typeFullAttentiondimself.config.qk_rope_head_dimnum_head1。二者虽然是不同的 cache 条目但attn_type相同因此会被归入同一个 manager 统一管理。其层级关系如下4.2 CacheUtilscache_utils.py提供以下工具函数定义位置cache_utils.pyvalidate_cache_info()校验 cache 元数据合法性。calculate_block_num()计算各 attention type 的 block 数量。allocate_cache_tensors()根据 block 规划结果分配底层 cache tensor。prepare_block_tables()根据各 request 的 block table 构造批次block_table。prepare_slot_mapping()根据position_ids和block_table生成slot_mapping。4.2.1 详解 calculate_block_num()calculate_block_num()会计算要给每种attn_type对应的SingleTypeKVCacheManager分配多少 block 数量。 它会先按attn_type将 cache 分成两类再分别计算block_num1. 固定 block 数的attn_type这类attn_type需要的 cache 数据不随着请求总长度线性增长例如SlidingWindow只需要窗口长度的 cache因此只给其分配固定大小的 block。代码中通过FIXED_BLOCK_ATTN_TYPES对这类attn_type进行标识。 当前典型例子是SlidingWindow。对其block_num的计算思路是先为单个请求估算“窗口内最多需要保留多少真实 block”再乘以最大并发数batch_size_per_dp_rank。具体计算为fixed_block_num max_concurrency * (ceil((sliding_window 2 * next_n) / block_size) 1) 1其中sliding_window 2 * next_n 是用于每个请求在一次迭代内最多会使用的 cache 大小2 * next_n 为 MTP 场景下主模型与小模型共同需要额外使用的 cache 空间1 个用于每个请求覆盖滑动窗口跨 block 边界时的额外保留需求最后的 1 个用于null_block占位。2. 非固定 block 数的attn_type这类attn_type需要的 cache 数据随着请求总长度线性增长例如FullAttentionMLA、MHA、GQA等全注意力都属于FullAttention。 在计算这类attn_type的所需 block 数时在离线推理模式下可以直接根据固定的最大长度计算block_num_per_request ceil(offline_max_len / block_size) total_block_num block_num_per_request * max_concurrency 1这里额外的1同样用于null_block占位。在在线推理模型下则先统计所有这类 cache 的per_token_bytes再用剩余显存预算换算出可支持的最大 token 数最后再折算为 block 数。伪代码如下per_token_bytes sum(cache.dim * cache.num_head * dtype_size(cache.dtype)) available_memory available_kv_memory * kv_cache_mem_ratio - fixed_block_memory_bytes max_tokens floor(available_memory / per_token_bytes) non_fixed_block_num floor(max_tokens / block_size) for attn_type in paged_attn_types: block_num_by_type[attn_type] non_fixed_block_num 1其中额外的1用于预留null_block。整体上calculate_block_num()的顺序是先为固定 block 类 cache 预留内存再用剩余内存估算非固定 block 类 cache 的 block 数这样可以避免 SWA 这类必须预留的 cache 挤占后续计算时出现不确定性。4.2.2 构建 block_table 与 slot_mappingprepare_block_tables()和prepare_slot_mapping()的输出均为字典结构键为attn_type值为对应的block_tables / slot_mapping张量即block_tables[attn_type] block_table_tensorslot_mapping[attn_type] slot_mapping_tensor这种形式与多 attention type 的设计一致推理运行时可以按attn_type取出对应的block_table和slot_mapping。5. 新模型适配 Cache 管理的 Checklist本文档提供新模型接入 KV Cache 管理框架的完整步骤清单涵盖必选和可选任务帮助开发者定位具体代码位置。5.1 分析模型 Attention 形态【必选】5.1.1 确定每个 cache 的attn_type在模型 attention 类中分析每个 cache 的特性并确定其归属类型支持的attn_type对应 Manager 类特性说明FullAttentionFullAttentionManager全注意力cache 随序列增长线性追加SlidingWindowSlidingWindowManager滑动窗口注意力支持过期块回收检查位置single_type_kv_cache_manager.py 的ATTN_TYPE_MANAGER_MAPATTN_TYPE_MANAGER_MAP: Dict[str, Type[SingleTypeKVCacheManager]] { FullAttention: FullAttentionManager, SlidingWindow: SlidingWindowManager, }判断逻辑若模型使用标准 MHA/MQA/GQA/MLA 等全注意力机制使用FullAttention若模型使用滑动窗口注意力如 gpt-oss使用SlidingWindow需额外设置sliding_window参数若现有ATTN_TYPE_MANAGER_MAP不满足需求参考步骤 3 新增 SingleTypeKVCacheManager 类5.2 在模型中配置 Cache 信息【必选】5.2.1 定义 cache 变量在 attention 类的__init__()中为每个 cache 定义占位 tensor用于存放每一层的 cache 数据示例代码位置modeling_deepseek.py self.nope_cache torch.Tensor([]) self.rope_cache torch.Tensor([])5.2.2 创建CacheEntry实例列表在 attention 类的__init__()中为每个 cache 创建CacheEntry实例并存入列表self.cache_entries。关键参数说明参见 cache_info.py参数说明cache_namecache 名称如k_cache、v_cache、nope_cacheattn_typeattention 类型从步骤 1 确定dim单个 attention head 的特征维度num_headcache 包含的 head 数量注意考虑 TP 切分后的实际 head 数dtypecache tensor 数据类型如torch.bfloat16、torch.int8needs_block是否参与 block 分配通常为Truetensor_setter回调函数用于将分配的 cache tensor 注册到 attention 类sliding_window仅SlidingWindow类型需要设置窗口大小示例代码位置modeling_deepseek.pyself.cache_entries [ CacheEntry( cache_namenope_cache, attn_typeself.attn_type, dimself.config.kv_lora_rank, num_head1, dtypedtype_nope, needs_blockTrue, tensor_setterlambda tensor, layerself: setattr(layer, nope_cache, tensor), ), CacheEntry( cache_namerope_cache, attn_typeself.attn_type, dimself.config.qk_rope_head_dim, num_head1, dtypedtype_rope, needs_blockTrue, tensor_setterlambda tensor, layerself: setattr(layer, rope_cache, tensor), ), ]5.2.3 遍历所有层构建ModelCacheInfo在模型类如DeepseekV3ForCausalLM中实现get_cache_info()方法。核心逻辑遍历self.model.layers对每个层取出CacheEntry实例列表layer.self_attn.cache_entries封装为LayerCacheInfo并加入列表将LayerCacheInfo列表封装为ModelCacheInfo实例并返回示例代码位置modeling_deepseek.pydef get_cache_info( self, ) - ModelCacheInfo: layers self.model.layers if not self.is_mtp else self.model.layers.values() layer_infos [] for layer_idx, layer in enumerate(layers): layer_infos.append( LayerCacheInfo( layer_idxlayer_idx, cacheslist(layer.self_attn.cache_entries), ) ) return ModelCacheInfo( num_layerslen(layer_infos), block_sizeself.block_size, layer_infoslayer_infos, )注意确保模型类持有block_size在模型类__init__()中保存block_size供get_cache_info()使用self.block_size infer_config.scheduler_config.block_size5.3 新增SingleTypeKVCacheManager子类【可选】仅在现有ATTN_TYPE_MANAGER_MAP无法满足新 cache 类型需求时执行。5.3.1 新建 Manager 子类位置single_type_kv_cache_manager.py需要实现的函数函数名说明get_num_blocks_to_allocate(request_id, num_tokens)计算本次需要新增多少 blockallocate_new_blocks(request_id, block_num_to_allocate, num_tokens)执行真实 block 分配更新req_to_blocksget_num_skipped_tokens(num_computed_tokens)计算可跳过分配的 token 数remove_skipped_blocks(request_id, total_computed_tokens)【可选】重写用于特殊回收机制validate_and_build_kwargs(group_entries)【静态方法】用于从CacheEntry提取和校验类型特定参数参考实现single_type_kv_cache_manager.py 的FullAttentionManager、SlidingWindowManager5.3.2 注册到ATTN_TYPE_MANAGER_MAP位置single_type_kv_cache_manager.pyATTN_TYPE_MANAGER_MAP: Dict[str, Type[SingleTypeKVCacheManager]] { FullAttention: FullAttentionManager, SlidingWindow: SlidingWindowManager, NewAttentionType: NewAttentionManager, # 新增 }5.4 扩展calculate_block_num()【可选】仅在新增 Manager 且现有 block 计算逻辑不满足需求时执行。calculate_block_num()会在 cache 管理初始化对每个attn_type分配 block 内存块如果新增了 Manager需要检查或新增分配逻辑。5.4.1 判断是否属于固定 block 类型判断新增的 Manager 属于哪一种类型固定 block 类型block 数不随序列长度增长如 SlidingWindow非固定类型block 数随序列长度线性增长如 FullAttention5.4.2 固定 block 类型扩展calculate_fixed_block_memory_bytes()若新增固定 block 类型首先在cache_utils.py 的FIXED_BLOCK_ATTN_TYPES中新增该attn_typeFIXED_BLOCK_ATTN_TYPES {SlidingWindow, NewFixedBlockType} # 新增然后需要在cache_utils.py 的calculate_fixed_block_memory_bytes()函数中补充计算逻辑if NewFixedBlockType in cache.attn_type: # 计算该类型的 fixed_block_num fixed_block_num ... else: raise AttributeError(...)5.4.3 确保calculate_block_num()正确处理位置cache_utils.py若新增非固定 block 类型检查现有的 block 计算逻辑是否满足新 Manager 需求。5.5 在 Attention forward 中接入block_table/slot_mapping【必选】5.5.1 从forward_metadata获取 block 信息block_table和slot_mapping存储在forward_metadata中均为Dict[str, torch.Tensor]结构生成位置cache_utils.pyprepare_block_tables()cache_utils.pyprepare_slot_mapping()数据结构block_tables forward_metadata.block_table # Dict[attn_type, Tensor] slot_mapping forward_metadata.slot_mapping # Dict[attn_type, Tensor]5.5.2 按attn_type取出对应 tensor在 attention forward 中使用self.attn_type作为 keyblock_table forward_metadata.block_table[self.attn_type] slot_mapping forward_metadata.slot_mapping[self.attn_type]注意事项同一层多个 cache如k_cache和v_cache共享相同的attn_type使用同一份block_table和slot_mapping若一层包含多种attn_type的 cache需分别取出对应的 tensor5.6 验证完整链路【必选】5.6.1 初始化链路验证以 execution_engine.py 的_init_cache_manager()函数为主线步骤函数位置验证内容1ModelWorker.get_cache_info()→ model_worker.py模型正确返回ModelCacheInfo2validate_cache_info()→ cache_utils.pycache 元数据合法3calculate_block_num()→ cache_utils.py各类型 block 数正确计算4allocate_cache_tensors()→ cache_utils.pycache tensor 正确分配5create_single_type_managers()→ single_type_kv_cache_manager.pyManager 正确创建6KVCacheManager(...)→ kv_cache_manager.py顶层 Manager 初始化完成5.6.2 运行期链路验证步骤函数位置验证内容1KVCacheManager.allocate_slots()→ kv_cache_manager.py能成功分配 block2prepare_block_tables()→ cache_utils.pyblock_table正确生成3prepare_slot_mapping()→ cache_utils.pyslot_mapping正确生成4模型前向能正确消费对应attn_type的 block 信息关键调用点Prefill 阶段scheduler.py 的_schedule_prefill_batch()调用allocate_slotsDecode 阶段scheduler.py 的_schedule_decode_batch()调用allocate_slotsEngine 构建execution_engine.py 的_build_model_inputs()调用prepare_block_tables和prepare_slot_mapping5.6.3 释放链路验证步骤函数位置验证内容1KVCacheManager.free()→ kv_cache_manager.py请求结束时正确释放 block2SingleTypeKVCacheManager.free()→ single_type_kv_cache_manager.py各类型 Manager 正确释放3BlockPool.free_blocks()→ block_pool.pyblock 正确回收无重复释放调用点scheduler.py 的_process_batch_output()在请求完成时调用kv_cache_manager.free(request.request_id)【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考