深入理解vLLM框架KV cache管理：逻辑层、物理层与PagedAttention详解！-开发者社区

引言

KV cache的管理是vLLM框架最关键内容之一，在框架升级到V1后其逻辑进行了一次大的调整。为更好的了解KV cache的管理逻辑，本文结合代码(v0.10.2版本)，从整体架构到关键细节进行讲解，涵盖逻辑层、物理层以及两者间的联系，帮助使用者/开发者对整体结构有个基本了解。

一、整体框架

1.1基本原理

KV cache管理逻辑以PagedAttention为基础进行构建，分了逻辑层与物理层，该方式类似于操作系统的虚拟内存（virtual memory）管理。虽然vLLM版本在快速迭代更新，但这个基础的逻辑保持一致，所以学习PagedAttention是了解KV cache管理的第一步。

PagedAttention的核心逻辑是将Attention运算中的KV值按照虚拟映射的方式管理起来，如下图所示。图中有两个请求request A和B，他们各自拥有自己的逻辑块（logical kv blocks），通过对应的映射表（block table）找到每个词在物理块（physical kv blocks）中的位置。

PagedAttention基本原理示意

这种方式优势：

能够充分利用显存，降低KV cache导致的存储碎片化问题；
减少物理显存的反复申请/释放操作，提升效率；

目前在V1版本中KV cache管理还融合了前缀树，更好地适配了Prefix cache功能。整体的架构如下图所示，逻辑层由KV Manager管理、物理层由Model Runner处理；通过Scheduler（调度器）作为信息传递的桥梁，衔接了逻辑层与物理层。定义的cache的管理单位包括：池（pool）表(table)、层(layer)、块(block)和槽(slot)。

slot为最小管理单元，每个token占一个slot；
block为请求分配的基本单位，一个block包含多个slot；
pool为逻辑层block的管理合集，通过链表将block数据组织起来；
table管理请求与数据的映射，一个table可包含多个请求的信息。位于物理层；
layer是物理层申请的数据单位，单个layer包含多个block被所有请求共用；

模块之间运行的关键步骤：

Scheduler分配资源给请求，通过KV Manager申请逻辑blocks；
KV Manager把Pool中空闲的blocks选中后给到对应请求；
分配好逻辑blocks后Scheduler构建scheduler.output传递给ModelRunner；
ModelRunner为每条请求创建block table，并生成slot_mapping；
计算时把slot_mapping传入attention，就能够从物理KV blocks上面找到所需的数据。

1.2cache的重复利用

KV cache manager中有个关键功能：重复利用已计算的KV cache，该功能可以降低重复前缀的计算。**使用cache对计算量的影响有多大？**从计算量的角度来分析一下这个问题：

在prefill阶段，仅考虑Attention（MLA）模块和FFN（MoE）模块的flops计算，通过增加已计算的前缀缓存（prefix cache）长度获得算量的变化。相关公式如下：

选择总长度不同的序列1k、4k、8k和16k分别计算得到如下曲线。通过对比可以知道：

KV cache匹配命中率越高计算量越小。
输入总序列越长，KV cache对计算的影响越大。

2.1blocks的管理逻辑

KV cache的存储是以block为基本单位组织的。据需求设定block size，表示block里面可存储的tokens的数量，用slot来表示token在block中位置。所有的block都在pool里面。

Block的定义：KVCacheBlock，它是一个双向链表。

# 代码位置：vllm/vllm/v1/core/kv_cache_utils.py # KVCacheBlock定义： class KVCacheBlock: """KV-cache block metadata.""" # Block ID, ranging from 0 to num_gpu_blocks - 1. block_id: int # Reference count. ref_cnt: int = 0 # The hash key (block hash + group id) of the block, only available # when the block is full and cached. _block_hash: Optional[BlockHashWithGroupId] = None # Used to construct a doubly linked list for free blocks. # These two attributes should only be manipulated by FreeKVCacheBlockQueue. prev_free_block: Optional["KVCacheBlock"] = None next_free_block: Optional["KVCacheBlock"] = None # Whether the block is a null block that should never be cached. is_null: bool = False

其中关键参数ref_cnt，表示有多少请求正在使用该block；_block_hash 用作已完成计算块的唯一标记。KVCacheBlock双向链表示意：

组织blocks的几个关键模块：

块池（Block Pool）：存储KVCacheBlock，一般初始化时决定数量，避免多次的python操作。
空闲队列（Free Block Queue）：空闲的块的队列，仅存储头尾节点指针信息。
缓存协调模块（KVCache Coordinator）：协调不同的KV cache组。

# 当前版本模块的定义位置： # BlockPool vllm/vllm/v1/core/block_pool.py # FreeKVCacheBlockQueue vllm/vllm/v1/core/kv_cache_utils.py # KVCache Coordinator vllm/vllm/v1/core/kv_cache_coordinator.py

2.2KV cache Manager的运行逻辑

KV cache Manager关键动作包括：开辟、释放、淘汰，逻辑流程如下图所示：

执行逻辑示意图（未考虑优先级）

a> 内存开辟

a.1 检查是否有足够的空间为新请求来开辟block；

a.2 从cache block检查是否可以复用，可以复用从free队列剔除，引用计数+1；

a.3 若无复用数据从free队列中头中弹出一个block用于写数据

a.4 如果block写满了数据，则被cached block标记

b> 内存释放

request使用完后，将block送回free队列里面，如果是cached block，引用计数-1，最后的这个block 最先放入队列中，它可能未被cache；

c> 淘汰策略

根据LRU策略淘汰队首的block cache，并从cached blocks字典中去除记录。

基本操作在近期KVCacheManager迭代中有所更改，以适配新特性比如PD connector、KV cache group等。具体细节参看代码：vllm/vllm/v1/core/kv_cache_manager.py

KV cache管理中比较重要的功能：

**a. Connector：**主要解决跨实例/节点之间传输问题，建议参看之前的总结，

b. Prefix cache：前缀匹配，在prefill阶段，重复利用已计算的KV cache，即把保存的KV cache数据（未被淘汰）提供给新请求使用。实现中，在创建blocks时检查是否有已计算的blocks，然后复用这些已计算的blocks。

prefix cache实现由scheduler和kv manager共同参与完成。参数“enable_caching”控制功能的开启与否，默认开启prefix cache。

主要执行动作：scheduler计算tokens时会触发匹配函数，然后交由kv manger进行匹配计算并返回已匹配的blocks；在为tokens申请新slots时，对block的使用情况进行标记。

prefix cache关键流程

这份完整版的大模型 AI 学习和面试资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】

三、应用层（物理层）

应用层承载实际的KV cache数据，在vLLM框架中也可其视为“物理层”。应用层通过torch创建一个连续的tensor实现，Attention运算时使用该tensor的部分数据。torch为作为应用层的基础，其有一套独立的显存管理逻辑，该逻辑依赖底层硬件API库。作为框架开发者/使用者，我们不需要太多关注底层逻辑，聚焦到torch之上的内容即可。

3.1kv cache的创建

kv cache的创建主要是两个步骤：

步骤1，**可用空间计算。**在初始化阶段，确认kv cache可用空间。
步骤2，**层的创建。**根据模型特点为每一层分配一个固定的空间。

可用空间计算

kv cache的可用空间是在预热(dummy run)阶段确定的。目前可用的空间（available_gpu_memory）计算有两种方式：

# vllm/vllm/v1/engine/core.py # _initialize_kv_caches函数： if os.environ.get("VLLM_ELASTIC_EP_SCALE_UP_LAUNCH") == "1": dp_group = getattr(self, "dp_group", None) assert dp_group is not None self.available_gpu_memory_for_kv_cache = \ ParallelConfig.sync_kv_cache_memory_size(dp_group, -1) available_gpu_memory = [ self.available_gpu_memory_for_kv_cache ] * len(kv_cache_specs) else: # Profiles the peak memory usage of the model to determine how # much memory can be allocated for kv cache. available_gpu_memory = ( self.model_executor.determine_available_memory()) self.available_gpu_memory_for_kv_cache = \ available_gpu_memory[0]

其中通过dummy run取峰值(peak)计算的方式的执行逻辑在determine_available_memory函数中完成。

# 代码位置：vllm/vllm/v1/worker/gpu_worker.py determine_available_memory # 主要函数：determine_available_memory

计算满足：KV blocks = 总显存 × gpu_memory_utilization −（模型显存 + 激活峰值 + 非 torch 显存）

blocks的大小计算如下，其中page size 跟模型结构相关，不同层的计算可以不一样。

# 代码位置：vllm/vllm/v1/core/kv_cache_utils.py def get_num_blocks(vllm_config: VllmConfig, num_layers: int, available_memory: int, page_size: int) -> int: """ Get the number of kv cache blocks. Args: vllm_config: The global VllmConfig num_layers: The number of layers available_memory: Memory available for KV cache in bytes. page_size: The page size of the KV cache. """ num_blocks = int(available_memory // page_size // num_layers) num_blocks = max(num_blocks, 0) num_blocks = may_override_num_blocks(vllm_config, num_blocks) return num_blocks # 其中同构的attention的page_size计算： # 代码位置：vllm/vllm/v1/kv_cache_interface.py def page_size_bytes(self) -> int: # For MLA we only store a single latent vector coef = 1 if self.use_mla else 2 return coef * self.block_size * self.num_kv_heads * self.head_size \ * get_dtype_size(self.dtype)

层的创建

blocks的物理显存分配是由pytorch完成，创建的连续tensor视为KV cache所能使用的物理显存。在GPU runner中有一个_allocate_kv_cache_tensors函数，通过torch.tensor逐层创建kv_cache_tensor。模型每层共用tensor数据，比如MLA模块会创建nope和rope模块，nope和rope可以创建在一起，也可以由两个tensor承载。

nope与rope连续

nope与rope分开

# 代码位置：vllm/v1/worker/gpu_model_runner.py kv_cache_raw_tensors: dict[str, torch.Tensor] = {} kv_cache_raw_tensors: dict[str, torch.Tensor] = {} for kv_cache_tensor in kv_cache_config.kv_cache_tensors: tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=self.device) for layer_name in kv_cache_tensor.shared_by: kv_cache_raw_tensors[layer_name] = tensor

3.2层数据分配

请求的逻辑blocks分配好后，通过scheduler.out传递给model runner，进而构造request与物理blocks的映射关系；scheduler.out里面还标记了每个请求中已计算完成的tokens。

在model runner中需要完成：scheduler.out信息要转化为PagedAttention所需要的KV数据。

先看一下PagedAttention运算时的数据形态，以FA为例，如下所示。根据Attention的一般计算可知，其输入是：QKV数据、历史的KV数据。

# Attention的输入 Args: query: shape = [num_tokens, num_heads, head_size] key: shape = [num_tokens, num_kv_heads, head_size] value: shape = [num_tokens, num_kv_heads, head_size] kv_cache: shape = [2, num_blocks, block_size, num_kv_heads, head_size] attn_metadata: Metadata for attention. # 输出： Returns: shape = [num_tokens, num_heads * head_size]

其中，query/key/value是当前forward产生的数据，kv_cache参数是整层的数据，该数据所有请求公用。

接着要明确问题是：Attention拿kv_cache的哪些数据？

在数据读取时仅有block的索引信息是不够的，还需要知道token在block中的具体位置，而这个信息存在attn_metadata的slot_mapping中(数据类型为torch.Tensor)。

slot_mapping如何运作？

slot_mapping存储了每个token在layer数据中的位置，它包含的元素总数与请求的tokens数量相等，其运算需要与block size结合。

计算方式：数据进行整除+取余运算，整数表示所在block id，余数表示在对应的block id里面的slot（槽位）。举个例子，假设block size=128，slot_mapping=[958, 714, 238, 427]，计算如下：

scheduler是多个请求同时下发，映射关系有多组，slot_mapping如何构造？

vLLM中通过映射表（block table）来记录每个请求信息，并处理多请求问题。用BlockTable类实现请求的增、删操作，代码如下（节选）：

# 代码位置： vllm/vllm/v1/worker/block_table.py class BlockTable: def __init__( self, block_size: int, max_num_reqs: int, max_num_blocks_per_req: int, max_num_batched_tokens: int, pin_memory: bool, device: torch.device, ): # .... self.block_table = self._make_buffer(max_num_reqs, max_num_blocks_per_req, dtype=torch.int32) self.num_blocks_per_row = np.zeros(max_num_reqs, dtype=np.int32) self.slot_mapping = self._make_buffer(self.max_num_batched_tokens, dtype=torch.int64) def append_row( self, block_ids: list[int], row_idx: int, ) -> None: if not block_ids: return num_blocks = len(block_ids) start = self.num_blocks_per_row[row_idx] self.num_blocks_per_row[row_idx] += num_blocks self.block_table.np[row_idx, start:start + num_blocks] = block_ids # 添加请求： def add_row(self, block_ids: list[int], row_idx: int) -> None: self.num_blocks_per_row[row_idx] = 0 self.append_row(block_ids, row_idx)

关键参数：**self.block_table，**其大小等于最大请求数乘以单个请求blocks数量最大值。它存储了所有需要运行请求相关的blocks信息。

block table通过add_row添加请求，其中row_idx参数是指传入的请求索引（req_index）。slot_mapping的生成计算如下（节选），由compute_slot_mapping函数完成。

# 代码位置： vllm/vllm/v1/worker/block_table.py def compute_slot_mapping(self, req_indices: np.ndarray, positions: np.ndarray) -> None: # ..... block_table_indices = (req_indices * self.max_num_blocks_per_req + positions // self.block_size) # 找到请求索引的token对应的 block位置 block_numbers = self.block_table.np.ravel()[block_table_indices] # 展平后取索引数据 block_offsets = positions % self.block_size # 转化为在每个block中的位置偏移 np.add(block_numbers * self.block_size, block_offsets, out=self.slot_mapping.np[:req_indices.shape[0]]) # 构造slot_mapping数据。

这个计算实际上就是一个二维转一维的运算。举个例子帮助理解，假设block_size大小为5，max_num_blocks_per_req大小为4，索引某个token的计算如下：

接下来，需要了解req_indices，positions是如何计算出来的？

代码在ModelRunner的_prepare_inputs函数里面，根据scheduler_output.num_scheduled_tokens的数据来构造req_indices与positions。

注释里面给了一个例子，[2, 5, 3] 表示的是有三个请求，每个元素表示请求需要的tokens数。其中num_computed_tokens_cpu参数是指已计算过的tokens数量。

# 代码位置：vllm/vllm/v1/worker/gpu_model_runner.py # def _prepare_inputs # ... req_ids = self.input_batch.req_ids tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] num_scheduled_tokens = np.array(tokens, dtype=np.int32) # 调度到每个请求上的数据 max_num_scheduled_tokens = max(tokens) # Get request indices. # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2] req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens) # cu_num_tokens: [2, 5, 3] -> [2, 7, 10] # arange: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2] cu_num_tokens, arange = self._get_cumsum_and_arange( num_scheduled_tokens) # Get positions. positions_np = self.positions.np[:total_num_scheduled_tokens] np.add(self.input_batch.num_computed_tokens_cpu[req_indices], arange, out=positions_np)

至此分析完了scheduler_output与slot_mapping的计算逻辑，还有两个细节点：

i> KV cache group

为了适配异构形态、或者不同Attention模块（MHA/MLA/GQA/slide等）的kv cache混合使用设计了kv cache group概念，对应的block size可以设置不同，用不同的block table来管理每个KV cache group的映射关系。不同KV cache group之间协同工作由KVCacheCoordinator管理。