模型推理框架vllm-3——KVCache管理器

📅 2026/7/1 1:59:40
模型推理框架vllm-3——KVCache管理器
conda create -n vllm_server python3.12conda activate vllm_serverpip install vllm0.17.1在之前介绍vllm的调度过程在token_budge下去处理waiting以及running队列对于里面如何对块进行管理没有介绍比如说在running和waiting队列中虽然在token_budege下running队列处理过程通过理论计算比如说我需要running生成一个token你的显存能够你去生成这一个token吗以及在waiting队列处理过程中获取我已经计算过的token以及prefill能不能处理这么多new_tokens这些都需要kv_cache_manager过程进行管理比如说下面代码过程PYTHONCopy# running 队列中获取 num_new_tokens 是不是能够被分配到blocknew_blocks self.kv_cache_manager.allocate_slots(request,num_new_tokens,num_lookahead_tokensself.num_lookahead_tokens,)# waiting 队列获取已经计算过的tokennew_computed_blocks, num_new_local_computed_tokens (self.kv_cache_manager.get_computed_blocks(request))# waiting 中prefill阶段能不能将所有的 num_new_tokens去分配blocknew_blocks self.kv_cache_manager.allocate_slots(request,num_new_tokens,num_new_computed_tokensnum_new_local_computed_tokens,new_computed_blocksnew_computed_blocks,num_lookahead_tokenseffective_lookahead_tokens,num_external_computed_tokensnum_external_computed_tokens,delay_cache_blocksload_kv_async,num_encoder_tokensnum_encoder_tokens,)在代码中kv_cache_manager初始化过程为PYTHONCopy# /vllm/v1/core/sched/scheduler.pyself.kv_cache_manager KVCacheManager(kv_cache_configkv_cache_config,max_model_lenself.max_model_len,enable_cachingself.cache_config.enable_prefix_caching,use_eagleself.use_eagle,log_statsself.log_stats,enable_kv_cache_eventsself.enable_kv_cache_events,dcp_world_sizeself.dcp_world_size,pcp_world_sizeself.pcp_world_size,hash_block_sizeself.block_size,metrics_collectorself.kv_metrics_collector,)对于初始化过程中几个关键参数解释如下1、kv_cache_config主要是管理 num_blocks设置vllm/v1/kv_cache_interface.py2、max_model_len模型支持的最大序列长度3、enable_caching开启前缀缓存。block分配整体过程在最开始介绍vllm中的整体框架中介绍如下几个参数1、block_szie这个一般就是默认16块也就是每一块block存储16个tokens2、预分配显存大小一般就是直接 设备显存x分配比率比如说24Gx0.9≈21.6G也就是vllm提前占用21.6G显存但是实际kv cache占用不一定就有21.6G在初始化过程中模型会进行一次forward去计算出模型运行期间除了 KV cache 以外所消耗的内存3、block数量int(available_memory // page_size // num_layers)page_szie一个 block在单层上的字节数计算过程为(KV)× block_size × num_kv_heads× head_size× dtype_bytes2 × 16 × num_kv_heads × head_size× dtype_bytes。num_layers所有kv_cache_groups中层数的最大值除此之外vllm为模型的每一层都分配了一个独立的、形状为 [2, num_blocks, block_size, num_kv_heads, head_size] 的KV缓存张量具体过程后面描述。那么vllm中整体分配过程如下1对于每一组输入 prompt在 prefill 阶段会根据 token 数提前分配 block⌈num_tokens / block_size⌉ 块例如 block_size165 token 只需 1 块。 prefill 阶段一次性把 prompt 的所有 token 的 KV 写入这些 block对于每一个block中几个参数block_id表示当前block的序号ref_cnt表示当前block被引用次数。 decode 阶段是增量式的每次只生成 1 个或少量新 token如果当前最后一个 block 还有空位就继续写入满了再申请新 block。 当 GPU KV cache 内存不足时调度器会根据策略默认 FCFS 优先 decode进行 preemption可能把等待中的请求换出移动到CPU中、丢弃重排或在某些配置下丢弃部分 running 请求。虽然输入是多组 prompt但 vLLM 会把它们拼接成一个“序列”来统一计算以最大化 GPU 利用率。每个序列的 position ids 从 0 独立开始通过 attention mask保证每组 token 只能看到自己组的信息跨组完全隔离。 在 decode 阶段vLLM 使用 slot_mapping tensor 来记录本次 forward 中每个要生成的新 token 对应到物理 KV cache 的哪个 slot 索引从而实现非连续 block 的高效寻址和写入。KVCacheManager处理过程除了KVCache在vllm中还有一个prefix cache其作用表示当多个请求中有相同的前缀时避免重复计算这部分内容不过值得注意的是必须是公共前缀中间相同并不能共享原因很简单decode阶段是用n-1去预测n如果两个序列中前k个都相同那么直接复用即可比如说下面例子中eg1你好帮我介绍武汉 eg2你好帮我介绍北京就可以直接复用 “你好帮我介绍” 这部分kv cacheeg1你好帮我介绍武汉 eg2你是一个旅游专家你好帮我介绍北京不能实现上面复用在KVCacheManagervllm/v1/core/kv_cache_manager.py中核心逻辑如下几个1、get_computed_blocks为当前的request找到他的prefix cache2、allocate_slots为当前的new_token去申请分配block 3、free释放所有的block当request被处理完之后 3、cache_blocks为prefix cache去分配block4、reset_prefix_cache清空所有的cache block。在了解核心逻辑之前了解block池创建过程。blockpool创建测试模型为Qwen/Qwen2-0.5B-Instruct后续具体数值和显存大小32G以及显存初始化大小0.9有关具体代码位置vllm/v1/core/block_pool.pykvcache block构建vllm/v1/core/kv_cache_utils.py在代码中vllm/v1/core/block_pool.py直接创建所有的blocksself.blocks: list[KVCacheBlock] [KVCacheBlock(idx) for idx in range(num_gpu_blocks)]其中 num_gpu_block大小为131928 而里面的KVCacheBlock创建过程比较简单为每一块block都去创建如下属性;值得注意的是每一个block都是一个双向队列因此prev以及next分别指向上下的block的idx而其它熟悉含义如下block_id当前block的序号、ref_cnt当前block被引用次数。值得注意的是上面过程只是创建了一个元数据对象还不知道具体的显存物理地址还是在代码vllm/v1/core/kv_cache_utils.py中的显存分配过程PYTHONCopy# vllm/v1/core/kv_cache_utils.py# 测试模型为 Qwen/Qwen2-0.5B-Instructdef get_kv_cache_config_from_groups(vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], available_memory: int):if len(kv_cache_groups) 1 and isinstance(kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs):...else:group_size max(len(group.layer_names) for group in kv_cache_groups) # 24page_size get_uniform_page_size([group.kv_cache_spec for group in kv_cache_groups]) # 8192assert group_size 0, group_size must be greater than 0num_blocks get_num_blocks(vllm_config, group_size, available_memory, page_size) # 131928kv_cache_tensors []for i in range(group_size):shared_by []for j in range(len(kv_cache_groups)):if i len(kv_cache_groups[j].layer_names):shared_by.append(kv_cache_groups[j].layer_names[i])kv_cache_tensors.append(KVCacheTensor(sizepage_size * num_blocks, shared_byshared_by))...kv_cache_groups对应的结果为测试模型为Qwen/Qwen2-0.5B-Instructkv_cache_groups[KVCacheGroupSpec(layer_names[model.layers.0.self_attn.attn, model.layers.1.self_attn.attn, ..., model.layers.23.self_attn.attn], kv_cache_specFullAttentionSpec(block_size16, num_kv_heads2, head_size64, dtypetorch.bfloat16, page_size_paddedNone, head_size_v64, sliding_windowNone, attention_chunk_sizeNone))]最终的返回内容为kv_cache_configKVCacheConfig(num_blocks131928, kv_cache_tensors[KVCacheTensor(size1080754176, shared_by[model.layers.0.self_attn.attn]), ...], kv_cache_groups[KVCacheGroupSpec(layer_names[model.layers.0.self_attn.attn, ...], kv_cache_specFullAttentionSpec(block_size16, num_kv_heads2, head_size64, dtypetorch.bfloat16, page_size_paddedNone, head_size_v64, sliding_windowNone, attention_chunk_sizeNone))])1080754176 page_size* num_blocks 8192x131928上述代码主要是计算需要分配的显存大小和block数量除此之外会对模型中每一层都去计算需要分配的kv_cache_tensor大小比如说上面代码计算得到每一层结果都是8192* 131928也就是我的就会对每一层分配一个这么大的内容去显存中占用直接去看具体显存分配过程在代码vllm/v1/worker/gpu_model_runner.py中的初始化kv_cache_tensor过程如下PYTHONCopy# vllm/v1/worker/gpu_model_runner.pydef initialize_kv_cache_tensors(self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]):...else:# 创建 kv_cache_tensorskv_cache_raw_tensors self._allocate_kv_cache_tensors(kv_cache_config)# 修改 kv_cache_tensor 形状kv_caches self._reshape_kv_cache_tensors(kv_cache_config, kv_cache_raw_tensors, kernel_block_sizes)return kv_caches对于里面的创建 kv_cache_tensors过程则是直接通过tensor torch.zeros(kv_cache_tensor.size, dtypetorch.int8, deviceself.device)去初始化一层的张量其大小为 kv_cache_tensorpage_size* num_blocks8192x131928而后将其分配给每一层for layer_name in kv_cache_tensor.shared_by: kv_cache_raw_tensors[layer_name] tensor而后将每层中分配得到的tensor大小通过 _reshape_kv_cache_tensors 处理将最开始的 page_size* num_blocks的修改为 kv_cache_shape让每个层真正需要的 KV cache 形状和 stride 布局其中得到的kv_cache_shape为(2, 131928, 16, 2, 64)对应(KV, num_blocks, block_size, num_kv_heads, head_size)。也就是将上面提到的 KVCacheTensor 中每个size大小等于KV× num_blocks× block_size× num_kv_heads× head_size× dtype_bytes这样一来整个 KV Cache 的物理显存占用和元数据管理由两大部分组成物理 KV 数据存储部分真正占大头显存 在 kv_cache_tensors 创建和 reshape 阶段vLLM 为模型的每一层 attention或每个 KV cache group提前分配好了对应的 KVCacheTensor。 形状为(2, num_blocks, block_size, num_kv_heads, head_size)比如说(2, 131928, 16, 2, 64) 这部分是真正存储 K 和 V 数值的显存。 每一层都有自己独立的或通过 view 共享底层内存的tensor。 所有层的 tensor 都使用相同的 num_blocks131928 作为第一维block 维度。全局 Block 元数据管理部分几乎不占显存 在 BlockPool 中一次性创建list[KVCacheBlock] [KVCacheBlock(idx) for idx in range(num_gpu_blocks)]这 num_gpu_blocks 个 KVCacheBlock 对象只保存元数据block_id、ref_cnt、block_hash、last_access_time 等不存储任何 KV 数据。总结block_pool想象成一个大池塘当一个请求需要分配新的block时就会从free_block_queue取出对应数量的block如果不够那就分配失败当请求完成释放所有block时就将请求下的所有block释放回free_block_queue。block分配逻辑上面简单介绍了如何去分配每层中的显存占用以及创建block下面介绍输入一个新的request如何去分配block进行占用。以上面waiting分配过程为例KVCacheManager中初始block协调器self.coordinator get_kv_cache_coordinator–get_kv_cache_coordinator去选择不同的协调器通过allocate_new_blocks去分配新的block其内部逻辑是通过 manager进行配分– manager实现逻辑在get_manager_for_kv_cache_spec中通过不同的attention计算方式选择不同的分配方式vllm/v1/core/single_type_kv_cache_manager.py。PYTHONCopy# get_kv_cache_coordinator 选择kvcache分配器 以KVCacheCooridantor为例 kv_cache_coordinator.pydef allocate_new_blocks(self,request_id: str,num_tokens: int,num_tokens_main_model: int,num_encoder_tokens: int 0,) - tuple[list[KVCacheBlock], ...]:return tuple(manager.allocate_new_blocks(request_id,num_encoder_tokensif isinstance(manager, CrossAttentionManager)else num_tokens,num_tokens_main_model,)for manager in self.single_type_managers)# 里面 self.single_type_managers 主要是通过 get_manager_for_kv_cache_spec获取其代码定义位置为 single_type_kv_cache_manager.py# get_manager_for_kv_cache_spec支持的attention分配逻辑spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] {FullAttentionSpec: FullAttentionManager,MLAAttentionSpec: FullAttentionManager,SlidingWindowSpec: SlidingWindowManager,ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,MambaSpec: MambaManager,CrossAttentionSpec: CrossAttentionManager,SinkFullAttentionSpec: SinkFullAttentionManager,}在get_manager_for_kv_cache_spec支持多种attention都是继承SingleTypeKVCacheManger主要去看其内部的allocate_new_blocks和allocate_new_computed_blocks这两部分代码逻辑在 new_blocks中比较简单直接计算需要分配的blocks数量cdiva//b用新的生成的tokens数量除block_size而后取整即可而后去我的block pool去拿出对应数量的block即可。对于 new_computed_blocks则是从prefix cache中去拿出已经被处理过程的cache那么上述整个cache block分配逻辑如下比如说我的模型只有4层而后我输入token数量是50个block_size是16这样一来就需要4个block才能存下所有的tokens首先去检查prefix cache也就是allocate_new_computed_blocks过程假设前 2 个逻辑 block 已经缓存在 KV cache 中那么req_to_blocks值假如就是block_table [10, 21, null, null]前两个是真实 hit 的 block_id后两个先占位而后通过allocate_new_blocks计算还需要2个block去存储剩下的tokens假设得到的是4567那么最后我的tokens存储的block_table就是[10, 21, 45, 67]那么去计算prefill处理这个过程就是直接拿着这个block_table去实际物理地址中去找上面我的每一层提前缓存了 (2, 131928, 16, 2, 64)–(2, num_blocks, block_size, num_kv_heads, head_size) 这么多tensor那么对于每一层PYTHONCopylayer-0: kv_cache_layer0[:, [10,21,45,67], ...]...layer-3: kv_cache_layer0[:, [10,21,45,67], ...]此时我的blok没有被填满而后持续deocde到被填满分配了4个刚刚好64个被填满此时需要新的block加入是89那么拿出添加到 block_table中即可而后依次类推。Prefix Cache逻辑Prefix Cache指的是缓存已处理请求的 kv-cache 块并在新请求到达时重用这些块前提是新请求与先前请求具有相同的前缀2。在推理框架SGLang中一个比较明显特点就是Prefix Cache不过两者在逻辑上存在差异vllm是基于block而SGLang则是基于token级别的也就是说vLLM中只有block被填满了斌且block内都是相同的才能复用prefix cache。在此之前在vllm中的block管理结构如下block都是一个双链表彼此之间都知道前后的block比如说第 0 个 blocktokens 0~15hash0 hash( NONE_HASH tuple(tokens[0:16]) extra_keys )第 1 个 blocktokens 16~31hash1 hash( hash0 tuple(tokens[16:32]) extra_keys )以此类推后续都是如此。不过必须注意的是只有在你的block被完全填满比如说有16个token之后才会为这个block添加上hash标志在介绍调度器中[对于waiting队列管理中有一个逻辑就是]( https://www.big-yellow-j.top/posts/2026/03/15/vllm-2.html#:~:text判断完毕之后- ,2%E3%80%81%E5%BE%97%E5%88%B0%E5%B7%B2%E7%BB%8F%E8%AE%A1%E7%AE%97%E7%9A%84tokens%E6%95%B0%E9%87%8F,-%EF%BC%8C%E9%A6%96%E5%85%88%E5%8E%BB%E5%88%A4%E6%96%AD)new_computed_blocks, num_new_local_computed_tokens self.kv_cache_manager.get_computed_blocks(request)通过kv_manager去获取已经被计算过的blocks以及新的本地计算计算的tokens这部分代码内容逻辑如下PYTHONCopy# kv_cache_manager.pydef get_computed_blocks(self, request: Request) - tuple[KVCacheBlocks, int]:if not self.enable_caching or request.skip_reading_prefix_cache:return self.empty_kv_cache_blocks, 0computed_blocks, num_new_computed_tokens (self.coordinator.find_longest_cache_hit(request.block_hashes, # 请求预先计算好的 block hash 列表链式request.num_tokens - 1))...return self.create_kv_cache_blocks(computed_blocks), num_new_computed_tokens从代码中对于被计算过的blocks获取方法直接通过find_longest_cache_hit去获取这部分代码内部逻辑是我直接去遍历所有的KVCacheBlock中去看哪些 hash 是对应的如果找到对应的那就可以直接复用即可[直接用官方例子]( https://docs.vllm.com.cn/en/latest/design/prefix_caching/#eviction-lru:~:text块哈希。- ,%E7%A4%BA%E4%BE%8B,-%C2%B6)来解释时间1缓存为空输入request0prompt为A-O总共19个token。直接从block pool中分配了 4 个块。其中 3 个已满那么就会被缓存直接使用SHA256进行编码编码逻辑是sha256( parent_hash 本块tokens extra )分别表示上一个列表、本block token、额外信息。第四个块部分填充了 3 个 token。时间2开始decode处理这样就会对 request0 使块3填满并请求新块以继续解码这个时候有4块被填满那么就会在cache block中建立4个hash时间 3request1 到达带有 14 个提示 tokenA-Jk-n其中前 10 个 token 与请求 0 相同。我们可以看到只有前 2 个块8 个 token命中缓存因为第 3 个块只匹配了 4 个 token 中的 2 个。其具体的命中缓存内部逻辑如下比如开始 request1在预分配的4个block其中前三块都是被填满建立索引假设是 *0 *1 *2 *3 在我的request0中建立hash是 *0 *1 *4 *5 *6那么会去调用find_longest_cache_hit(request.block_hashes)去cache blocks进行查找发现request2中 *2没有命中那么直接复用 *0 以及 *1。也就对应下面图像中 request1服用了ID0和ID1比如说代码PYTHONCopyclassmethoddef find_longest_cache_hit(cls, block_hashes, max_length, ...):computed_blocks tuple([] for _ in range(len(kv_cache_group_ids)))max_num_blocks max_length // block_sizefor block_hash in itertools.islice(block_hashes, max_num_blocks):# 直接去 blocks_hashes 中查找if cached_block : block_pool.get_cached_block(block_hash, kv_cache_group_ids):# 遇到匹配的cahcefor computed, cached in zip(computed_blocks, cached_block):computed.append(cached)else:breakif use_eagle and computed_blocks[0]:for computed in computed_blocks:computed.pop()return computed_blocks时间 4请求 0 完成并释放。块 2、3 和 4 按反向顺序添加到空闲队列但块 2 和 3 仍被缓存。块 0 和 1 未添加到空闲队列因为它们正在被请求 1 使用时间 5请求 1 完成并释放。时间 6请求 2 到达带有 29 个提示 token其中前 12 个 token 与请求 0 相同。请注意即使空闲队列中的块顺序是 7 - 8 - 9 - 4 - 3 - 2 - 6 - 5 - 1 - 0缓存命中的块即 0、1、2在分配前会被触碰并从队列中移除因此空闲队列变为 7 - 8 - 9 - 4 - 3 - 6 - 5。结果分配的块是 0已缓存、1已缓存、2已缓存、7、8、9、4、3已驱逐。总结KVCache中几个核心的代码文件1、kv_cache_manager.py对block进行分配不过核心调度代码在kv_cache_coordinator.py中包括输入新的token去计算需要多少block、prefix cache使用2、block_pool.py创建blook池核心就是根据计算得到的num_blocks去计算分配多少全局管理的blockKVCacheBlock3、gpu_model_runner.py直接为模型的每一层去占用物理显存并且对于占用的物理显存通过全局的KVCacheBlock进行管理。对于block pool想象成一个池子总共有num_blocks个每个池子只能存储block_size个token我提前对我模型中每一层都分配了一个(KV, num_blocks, block_size, num_kv_heads, head_size)大小的张量进行占用所有物理块共享这块连续的 GPU 内存通过block_id进行分页索引在新的prompt输入时候计算需要分配池子数量 len(tokens)//block_size而后会将所有的token都放到池子中对于填满了的池子会用一个hash标记放到block hash中新的输入时候进行hash计算如果在block hash中有重复的直接复用prefix cache过程而在具体计算中就直接根据block需要去分配的层中进行索引比如说某个prompt分配了[1,2]那么计算过程就是