vllm: kv cache

参考文献

vllm官方prefix缓存设计文档：https://github.com/vllm-project/vllm/blob/main/docs/design/prefix_caching.md

数据表现

原始计算轮数据：

FlashAttentionMetadata: (num_actual_tokens=997, max_query_len=997, query_start_loc=tensor([0, 997]), max_seq_len=997, seq_lens=tensor([997]), slot_mapping=tensor([16, 17, 18, 19, 20, ... , 1012)], block_table=tensor([1, 2, 3, 4, 5, 6, ... ,61, 62, 63, 0, ... , 0)])

BatchDescriptor: (num_tokens=8)

缓存轮数据：

FlashAttentionMetadata: (num_actual_tokens=5, max_query_len=5, query_start_loc=tensor([0, 5]), max_seq_len=997, seq_lens=tensor([997]), slot_mapping=tensor([2032, 2033, 2034, 2035, 2036, -1, -1, -1]), block_table=tensor([1, 2, 3, 4, 5, 6, ... , 61, 62, 127, 0, ... , 0)])

原理&概要设计

prefix使用block内存储的token ids作为hash的对象，以block为单位存储每个block的hash值，作为kv cache复用的检索对象。对于多模态数据而言，由于多模态数据使用的是相同token id的多模态占位符，所以附加embeding过程产生的多模态data hash值作为附加哈希值。

管理与复用

以Block为单位，对block内的token进行hash，通过hash进行管理和复用；

参与hash值计算的为三部分：

父哈希值：父哈希块的哈希值。

块内****词元：该块中包含的词元元组。包含精确词元的目的是降低潜在的哈希值碰撞风险。

附加哈希：用于确保该块唯一性的其他必要值，例如LoRA标识符、多模态输入哈希值（参见下方示例），以及用于在多租户环境中隔离缓存的随机盐值。

在多模态场景中，由于块内词元是mm_placeholder，为了区分图片，附加哈希会添加上多模态输入哈希值（如图片转image token时创建的hash值）；

数据结构

vLLM v1 中的前缀缓存在 KV cache manager中实现。其基本构建块是 "Block" 数据类（简化版）：

class KVCacheBlock:
    # The block ID (immutable)
    block_id: int
    # The block hash (will be assigned when the block is full,
    # and will be reset when the block is evicted).
    block_hash: BlockHash
    # The number of requests using this block now.
    ref_cnt: int
    # The pointers to form a doubly linked list for the free queue.
    prev_free_block: "KVCacheBlock | None" = None
    next_free_block: "KVCacheBlock | None" = None

image

详细设计

初始化

kv_cache由Scheduler创建和管理：self.kv_cache_manager = KVCacheManager

Block Allocation

1 New request

scheduler为新请求分配KV缓存块的工作流程：

scheduler调用kv_cache_manager.get_computed_blocks()来获取已计算完成的块序列。这是通过对request中的提示词进行哈希并查找缓存块来实现的。

scheduler调用kv_cache_manager.allocate_slots()。该函数执行以下步骤：

  1. 计算所需的新block数量，如果可用block不足则直接返回。

  2. "标记"已计算的block。将已计算block的引用计数加一，如果该block未被其他请求使用，则将其从空闲队列中移除。这是为了避免这些已计算的block被淘汰。具体示例见下一节说明。

  3. 从空闲队列头部弹出块来分配新block。如果弹出的头block是已缓存的block，则会"淘汰"该block，从现在起其他请求将无法再复用该block。

  4. 如果分配的block已填满token，我们会立即将其加入缓存block，使该block能在同一批次中被其他request复用。

2 Running request

scheduler为正在运行的请求分配KV缓存块的工作流程：

  1. scheduler调用kv_cache_manager.allocate_slots()。该函数执行以下步骤：

    1. 计算所需的新block数量，如果可用block不足则直接返回。

    2. 通过从空闲队列头部弹出块block分配新block。如果弹出的头block是已缓存的block，则会“淘汰”该block，从现在起其他请求将无法再复用该块block

    3. 将token id追加到现有块以及新块的插槽中。如果一个block已满，则将其加入缓存块进行缓存。

Eviction (LRU)

当空闲队列的头block（最近最少使用block）已被缓存时，我们必须淘汰该block以防止其被其他reques使用。具体而言，淘汰过程包含以下步骤：

  1. 从空闲队列头部弹出该block。这是即将被淘汰的LRU block。

  2. 从cache blocks中移除该block的ID。

  3. 移除该block的哈希值。

示例

在本示例中，我们假设block大小为4（每个block可缓存4个token），且KV缓存管理器中共有10个block。

时刻1：缓存为空，新请求到达。我们分配了4个块。其中3个块已满并被缓存。第四个块部分填充了3个令牌（容量为4）。

image

时刻2：请求0使第3块填满，并要求分配新块以继续解码。我们缓存第3块并分配第4块。

image

时刻3：请求1到达，包含14个提示令牌，其中前10个令牌与请求0相同。我们可以看到只有前2个块（8个令牌）命中缓存，因为第3个块只匹配了4个令牌中的2个。

image

时刻4：请求0完成并被释放。块2、3和4按逆序加入空闲队列（但块2和3仍处于缓存状态）。块0和1未被加入空闲队列，因为它们正被请求1使用。

image

时刻5：请求1完成并被释放。

image

时刻6：请求2到达，包含29个提示令牌，其中前12个令牌与请求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（被淘汰）。

image

代码实现

数据结构

1）block_size: 每个kv cache block的slots数量

slots数量：即一个block能够存放多少token

默认情况下，每个KV Cache Block包含16个slots（tokens）。可选值：有效的block size包括 1, 8, 16, 32, 64, 128, 256；标准CUDA attention后端只支持最大32的block size，更大的size（64+）仅用于MLA（Multi-head Latent Attention）后端。

2）block_num: kv cache block数量

vllm执行虚拟前向传播（Dummy Forward Pass），使用最大batch size和最大序列长度来模拟峰值内存使用。可以分配给kv blocks的内存由(memory*utilization-model_weights-peak_usage)计算得到。实现细节而言，使用max_model_len来计算单个request推理需要的最大显存，max_num_batched_tokens来计算

max_model_len: 限制单个请求中"输入tokens + 输出tokens"的总和上限。

max_num_seqs: 控制vLLM能够同时处理的最大独立请求（序列）数量；默认值128。

max_num_batched_tokens: 限制单个调度步骤(schedule step)中所有请求的token总数预算，在调度器中通过SchedulingBudget类管理；默认值2048。

实验示例：

Model=Qwen3-vl-2B【4.24 GiB】, gpu_memory_utilization=0.9, 一张图片(from采集数据)的token占用约1000 tokens。

max_model_len=4096，max_num_seqs=5，max_num_batched_tokens=2048；Available KV cache memory: 16.87 GiB，GPU KV cache size: 157,952 tokens；

max_model_len=8192，max_num_seqs=5，max_num_batched_tokens=2048；Available KV cache memory: 16.87 GiB，GPU KV cache size: 157,952 tokens；

max_model_len=4096，max_num_seqs=10，max_num_batched_tokens=2048；Available KV cache memory: 16.87 GiB，GPU KV cache size: 157,952 tokens；

max_model_len=4096，max_num_seqs=5，max_num_batched_tokens=4096；Available KV cache memory: 16.59 GiB，GPU KV cache size: 155,296 tokens；

结论：模型大小，gpu_memory_utilization为KV cache能够容纳tokens数量的主要影响项，max_num_batched_tokens为轻微影响项。

初始化

GPUModelRunner.initialize_kv_cache_tensors

创建kv_caches: [torch.Size([2, 8243, 16, 8, 128])] * layer_num；注：8243为block数量，16为block大小；

将kv_caches与self.compilation_config.static_forward_context, self.kv_caches绑定；

计算过程中

cache 使用的变量

attn_metadata

组织形式：dict，以模型层名为dict_key，每一个item是模型对应层使用的attn_metadata。

以两个req为例，拆解其中的重要变量如下：

1）seq_lens: 当前批次中每个req已经推理的token数量。

# req1 prefill
seq_lens=tensor([997])
# req1 decode，req2 prefill
seq_lens=tensor([998, 1001])

2）block_table: 当前推理需要访问的block

# req1 prefill
block_table = tensor([[1, 2, 3, 4, 5, 6, ...]])
# req1 decode，req2 prefill
block_table = tensor([[1, 2, 3, 4, 5, 6, ...],[64, 65, 66, 67, 68, 69, 70, 71, ...]])
  1. slot_mappings: 当前批次中每个 token 的 KV 值应该写入（或读取）物理显存的具体位置（slot）。
# req1 prefill（len=997）
slot_mappings = tensor([16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27,...，1012])
# req1 decode，req2 prefill（len=1+1001）
slot_mappings = tensor([1013, 1024, 1025,  ..., 2022, 2023, 2024])

变量的传递和使用

  1. set_forward_context

additional_kwargs = current_platform.set_additional_forward_context(

attn_metadata=attn_metadata,

vllm_config=vllm_config,

num_tokens=num_tokens,

)

forward_context = create_forward_context(

attn_metadata,

vllm_config,

slot_mapping,

additional_kwargs,

)

forward_context通过全局变量_forward_context（forward_context.py）传递给模型运行时

  1. 获取使用

forward_context = get_forward_context()

attn_metadata: AttentionMetadata = forward_context.attn_metadata

在 vllm/attention/layer.py 中，Attention.forward() 方法直接调用 get_forward_context

prefix cache

以单张图的prefix cache复用，展示复用的变量配置

  1. 不使用缓存：
# prefill
num_actual_tokens=997
block_table = tensor([[1, 2, 3, ..., 61，62，63]])
slot_mappings = tensor([16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27,...，1011])
seq_lens=tensor([997])
  1. 使用缓存：
# prefill with prefix cache
num_actual_tokens=5
block_table = tensor([[1, 2, 3, ..., 61，62，64]])
slot_mappings = tensor([1024, 1025, 1026, 1027, 1028, -1, -1, -1])
seq_lens=tensor([997])

注：3个-1，代表的是3个mm占位符token：<|vision_start|>{placeholder}<|vision_end|>；

  1. 缓存轮后续推理
# decode1
num_actual_tokens=1
block_table = tensor([[1, 2, 3, ..., 61，62，64]])
slot_mappings = tensor([1029])
seq_lens=tensor([998])
# decode2
num_actual_tokens=1
block_table = tensor([[1, 2, 3, ..., 61，62，64]])
slot_mappings = tensor([1030])
seq_lens=tensor([999])
# decode3
num_actual_tokens=1
block_table = tensor([[1, 2, 3, ..., 61，62，64]])
slot_mappings = tensor([1031])
seq_lens=tensor([1000])
posted on 2026-02-25 20:07  uestc001  阅读(245)  评论(0)    收藏  举报

博客园  ©  2004-2026
浙公网安备 33010602011771号 浙ICP备2021040463号-3