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

在之前的Blog分别介绍了vllm整体框架及使用以及在vllm中生成以及调度过程 ，本文主要介绍块管理器的原理（全部是基于：Version: 0.17.1）。

conda create -n vllm_server python=3.12
conda activate vllm_server
pip install vllm==0.17.1

在之前介绍vllm的调度过程在token_budge下去处理waiting以及running队列，对于里面如何对块进行管理没有介绍，比如说在running和waiting队列中，虽然在token_budege下running队列处理过程，通过理论计算比如说我需要running生成一个token，你的显存能够你去生成这一个token吗？以及在waiting队列处理过程中获取我已经计算过的token以及prefill能不能处理这么多new_tokens？这些都需要 kv_cache_manager过程进行管理，比如说下面代码过程：

# running 队列中获取 num_new_tokens 是不是能够被分配到block
new_blocks = self.kv_cache_manager.allocate_slots(
                        request,
                        num_new_tokens,
                        num_lookahead_tokens=self.num_lookahead_tokens,
                    )
# waiting 队列获取已经计算过的token
new_computed_blocks, num_new_local_computed_tokens = (
    self.kv_cache_manager.get_computed_blocks(request)
)
# waiting 中prefill阶段能不能将所有的 num_new_tokens去分配block
new_blocks = self.kv_cache_manager.allocate_slots(
    request,
    num_new_tokens,
    num_new_computed_tokens=num_new_local_computed_tokens,
    new_computed_blocks=new_computed_blocks,
    num_lookahead_tokens=effective_lookahead_tokens,
    num_external_computed_tokens=num_external_computed_tokens,
    delay_cache_blocks=load_kv_async,
    num_encoder_tokens=num_encoder_tokens,
)

在代码中 kv_cache_manager初始化过程为：

# /vllm/v1/core/sched/scheduler.py
self.kv_cache_manager = KVCacheManager(
    kv_cache_config=kv_cache_config,
    max_model_len=self.max_model_len,
    enable_caching=self.cache_config.enable_prefix_caching,
    use_eagle=self.use_eagle,
    log_stats=self.log_stats,
    enable_kv_cache_events=self.enable_kv_cache_events,
    dcp_world_size=self.dcp_world_size,
    pcp_world_size=self.pcp_world_size,
    hash_block_size=self.block_size,
    metrics_collector=self.kv_metrics_collector,
)

对于初始化过程中几个关键参数解释如下：1、kv_cache_config：主要是管理 num_blocks设置（vllm/v1/kv_cache_interface.py）；2、max_model_len：模型支持的最大序列长度；3、enable_caching：开启前缀缓存。

block分配整体过程

在最开始介绍vllm中的整体框架中介绍如下几个参数：1、block_szie：这个一般就是默认16块也就是每一块block存储16个tokens；2、预分配显存大小：一般就是直接 设备显存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在单层上的字节数计算过程为(K+V)$\times$ block_size $\times$ num_kv_heads$\times$ head_size$\times$ dtype_bytes=2 $\times$ 16 $\times$ num_kv_heads $\times$ head_size$\times$ dtype_bytes。num_layers：所有kv_cache_groups中层数的最大值，除此之外vllm为模型的每一层都分配了一个独立的、形状为 [2, num_blocks, block_size, num_kv_heads, head_size] 的KV缓存张量（具体过程后面描述）。那么vllm中整体分配过程如下¹：

对于每一组输入 prompt，在 prefill 阶段会根据 token 数提前分配 block：⌈num_tokens / block_size⌉ 块（例如 block_size=16，5 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 cache
eg1：你好，帮我介绍武汉？ eg2：你是一个旅游专家，你好，帮我介绍北京？==>不能实现上面复用

在KVCacheManager（vllm/v1/core/kv_cache_manager.py）中核心逻辑如下几个，1、get_computed_blocks：为当前的request找到他的prefix cache；2、allocate_slots：为当前的new_token去申请分配block， 3、free：释放所有的block当request被处理完之后 ；3、cache_blocks：为prefix cache去分配block；4、reset_prefix_cache：清空所有的cache block。在了解核心逻辑之前了解block池创建过程。

blockpool创建

测试模型为：Qwen/Qwen2-0.5B-Instruct（后续具体数值和显存大小（32G）以及显存初始化大小（0.9）有关）
具体代码位置：vllm/v1/core/block_pool.py
kvcache block构建：vllm/v1/core/kv_cache_utils.py

在代码中（vllm/v1/core/block_pool.py）直接创建所有的blocks： self.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 中的显存分配过程：

# vllm/v1/core/kv_cache_utils.py
# 测试模型为 Qwen/Qwen2-0.5B-Instruct
def 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) # 24
        page_size = get_uniform_page_size(
            [group.kv_cache_spec for group in kv_cache_groups]
        ) # 8192
        assert group_size > 0, "group_size must be greater than 0"
        num_blocks = get_num_blocks(
            vllm_config, group_size, available_memory, page_size
        ) # 131928
        kv_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(size=page_size * num_blocks, shared_by=shared_by)
            )
    ...

kv_cache_groups对应的结果为（测试模型为 Qwen/Qwen2-0.5B-Instruct ）： kv_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_spec=FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=64, dtype=torch.bfloat16, page_size_padded=None, head_size_v=64, sliding_window=None, attention_chunk_size=None))]
最终的返回内容为： kv_cache_config=KVCacheConfig(num_blocks=131928, kv_cache_tensors=[KVCacheTensor(size=1080754176, shared_by=['model.layers.0.self_attn.attn']), ...], kv_cache_groups=[KVCacheGroupSpec(layer_names=['model.layers.0.self_attn.attn', ...], kv_cache_spec=FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=64, dtype=torch.bfloat16, page_size_padded=None, head_size_v=64, sliding_window=None, attention_chunk_size=None))])
1080754176= page_size* num_blocks= 8192x131928

上述代码主要是计算需要分配的显存大小和block数量，除此之外会对模型中每一层都去计算需要分配的kv_cache_tensor大小，比如说上面代码计算得到每一层结果都是：8192* 131928，也就是我的就会对每一层分配一个这么大的内容去显存中占用，直接去看具体显存分配过程，在代码（vllm/v1/worker/gpu_model_runner.py）中的初始化kv_cache_tensor过程如下：

# vllm/v1/worker/gpu_model_runner.py
def initialize_kv_cache_tensors(self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]):
    ...
    else:
        # 创建 kv_cache_tensors
        kv_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, dtype=torch.int8, device=self.device)去初始化一层的张量其大小为 kv_cache_tensor（page_size* num_blocks=8192x131928），而后将其分配给每一层 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)对应 (K+V, num_blocks, block_size, num_kv_heads, head_size)。

也就是将上面提到的 KVCacheTensor 中每个size大小等于K+V$\times$ num_blocks$\times$ block_size$\times$ num_kv_heads$\times$ head_size$\times$ dtype_bytes

这样一来 物理显存占用了（创建 kv_cache_tensors过程）+block序号（直接创建所有的blocks），其中前者是每一层attention中（又多个block组成）分配得到的物理显存大小，后者则是表示每个block的属性（如block id等），对应过程比较简单，对于每一层attn都提前分配好了KVCacheTensor（(K+V, num_blocks, block_size, num_kv_heads, head_size)），在block序号创建过程中其数量恰为 num_blocks也就是说block序号就对应这一层中的所有内容

参考