引言

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.2 cache的重复利用

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

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

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

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

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

二、逻辑层

2.1 blocks的管理逻辑

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.2 KV 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关键流程

相关步骤代码：

#  scheduer 中调用kv manager
#  vllm/vllm/v1/core/sched/scheduler.py
#  触发匹配计算
# Get already-cached tokens.
if request.num_computed_tokens == 0:
# Get locally-cached tokens.
new_computed_blocks, num_new_local_computed_tokens = \
self.kv_cache_manager.get_computed_blocks(
request)
# vllm/vllm/v1/core/kv_cache_manager.py
# 进行命中率计算：
max_cache_hit_length = request.num_tokens - 1
computed_blocks, num_new_computed_tokens = (
self.coordinator.find_longest_cache_hit(request.block_hashes,
max_cache_hit_length))
#  vllm/vllm/v1/core/sched/scheduler.py
# scheduler申请slots：
new_blocks = self.kv_cache_manager.allocate_slots(
request,
num_new_tokens + num_external_computed_tokens,
num_new_local_computed_tokens,
new_computed_blocks,
num_lookahead_tokens=effective_lookahead_tokens,
delay_cache_blocks=load_kv_async,
num_encoder_tokens=num_encoder_tokens,
)
# vllm/vllm/v1/core/kv_cache_manager.py
# manger标记匹配的blocks：
# Touch the computed blocks to make sure they won't be evicted.
if self.enable_caching:
self.block_pool.touch(new_computed_block_list)
# vllm/vllm/v1/core/block_pool.py
# pool里面管理的block的引用计数（ref_cnt ）+1：
def touch(self, blocks: tuple[list[KVCacheBlock], ...]) -> None:
for blocks_per_group in blocks:
for block in blocks_per_group:
# ref_cnt=0 means this block is in the free list (i.e. eviction
# candidate), so remove it.
if block.ref_cnt == 0 and not block.is_null:
self.free_block_queue.remove(block)
block.ref_cnt += 1
# vllm/vllm/v1/core/single_type_kv_cache_manager.py
# pool释放请求对应的blocks：
def free(self, request_id: str) -> None:
"""
Free the blocks for the request.
Args:
request_id: The request ID.
"""
# Default to [] in case a request is freed (aborted) before alloc.
req_blocks = self.req_to_blocks.pop(request_id, [])
# Free blocks in reverse order so that the tail blocks are
# freed first.
# 此处有个按照逆序来淘汰策略，但该方式目前看非最佳（可留言讨论@kaiyuan）
ordered_blocks = reversed(req_blocks)
self.block_pool.free_blocks(ordered_blocks)
self.num_cached_block.pop(request_id, None)
# vllm/vllm/v1/core/block_pool.py
# 释放时，引用计数-1：
def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
# Materialize the iterable to allow multiple passes.
blocks_list = list(ordered_blocks)
for block in blocks_list:
block.ref_cnt -= 1
self.free_block_queue.append_n([
block for block in blocks_list
if block.ref_cnt == 0 and not block.is_null
])

三、 应用层（物理层）

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

3.1 kv 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管理。

相关的PR参考：

• https://github.com/vllm-project/vllm/pull/13296
• https://github.com/vllm-project/vllm/pull/25101

ii> DCP（Decode Context Parallel）

为了减少KV cache的冗余存储，开启Attention序列并行时，可根据CP的数量让不同设备存储KV cache的部分数据。

当前计算逻辑单个token的存储位置满足: rank_id = token_idx % cp_world_size。

---

关键细节介绍到这里，其它内容参考源码:

https://github.com/vllm-project/vllm/tree/releases/v0.10.2

新模型在不断推出，KV cache管理也在持续迭代。下一步的方向:

• 模型的适配: 当模型中Attention结构出现改进时，目前代码修改涉及逻辑层和应用层，这种方式很显然是高成本的，所以需要一种更加解耦降低修改量。
• 池化应用: 本地池化和跨节点的池化功能应该成为一种标配。
• 传输优化: 全掩盖的传输方式、不影响指标(TTFT/ITL）的传输方式的探索。

四、如何学习AI大模型？

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

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

第一阶段： 从大模型系统设计入手，讲解大模型的主要方法；

第二阶段： 在通过大模型提示词工程从Prompts角度入手更好发挥模型的作用；

第三阶段： 大模型平台应用开发借助阿里云PAI平台构建电商领域虚拟试衣系统；

第四阶段： 大模型知识库应用开发以LangChain框架为例，构建物流行业咨询智能问答系统；

第五阶段： 大模型微调开发借助以大健康、新零售、新媒体领域构建适合当前领域大模型；

第六阶段： 以SD多模态大模型为主，搭建了文生图小程序案例；

第七阶段： 以大模型平台应用与开发为主，通过星火大模型，文心大模型等成熟大模型构建大模型行业应用。

👉学会后的收获：👈

• 基于大模型全栈工程实现（前端、后端、产品经理、设计、数据分析等），通过这门课可获得不同能力；

• 能够利用大模型解决相关实际项目需求： 大数据时代，越来越多的企业和机构需要处理海量数据，利用大模型技术可以更好地处理这些数据，提高数据分析和决策的准确性。因此，掌握大模型应用开发技能，可以让程序员更好地应对实际项目需求；

• 基于大模型和企业数据AI应用开发，实现大模型理论、掌握GPU算力、硬件、LangChain开发框架和项目实战技能， 学会Fine-tuning垂直训练大模型（数据准备、数据蒸馏、大模型部署）一站式掌握；

• 能够完成时下热门大模型垂直领域模型训练能力，提高程序员的编码能力： 大模型应用开发需要掌握机器学习算法、深度学习框架等技术，这些技术的掌握可以提高程序员的编码能力和分析能力，让程序员更加熟练地编写高质量的代码。

1.AI大模型学习路线图
2.100套AI大模型商业化落地方案
3.100集大模型视频教程
4.200本大模型PDF书籍
5.LLM面试题合集
6.AI产品经理资源合集

👉获取方式：
😝有需要的小伙伴，可以保存图片到wx扫描二v码免费领取【保证100%免费】🆓