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" = Nonenext_free_block: "KVCacheBlock | None" = None
详细设计
初始化
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()。该函数执行以下步骤：
-
计算所需的新block数量，如果可用block不足则直接返回。
-
"标记"已计算的block。将已计算block的引用计数加一，如果该block未被其他请求使用，则将其从空闲队列中移除。这是为了避免这些已计算的block被淘汰。具体示例见下一节说明。
-
从空闲队列头部弹出块来分配新block。如果弹出的头block是已缓存的block，则会"淘汰"该block，从现在起其他请求将无法再复用该block。
-
如果分配的block已填满token，我们会立即将其加入缓存block，使该block能在同一批次中被其他request复用。
2 Running request
scheduler为正在运行的请求分配KV缓存块的工作流程：
-
scheduler调用
kv_cache_manager.allocate_slots()。该函数执行以下步骤：-
计算所需的新block数量，如果可用block不足则直接返回。
-
通过从空闲队列头部弹出块block分配新block。如果弹出的头block是已缓存的block，则会“淘汰”该block，从现在起其他请求将无法再复用该块block
-
将token id追加到现有块以及新块的插槽中。如果一个block已满，则将其加入缓存块进行缓存。
-
Eviction (LRU)
当空闲队列的头block（最近最少使用block）已被缓存时，我们必须淘汰该block以防止其被其他reques使用。具体而言，淘汰过程包含以下步骤：
-
从空闲队列头部弹出该block。这是即将被淘汰的LRU block。
-
从cache blocks中移除该block的ID。
-
移除该block的哈希值。
示例
在本示例中，我们假设block大小为4（每个block可缓存4个token），且KV缓存管理器中共有10个block。
时刻1：缓存为空，新请求到达。我们分配了4个块。其中3个块已满并被缓存。第四个块部分填充了3个令牌（容量为4）。
时刻2：请求0使第3块填满，并要求分配新块以继续解码。我们缓存第3块并分配第4块。
时刻3：请求1到达，包含14个提示令牌，其中前10个令牌与请求0相同。我们可以看到只有前2个块（8个令牌）命中缓存，因为第3个块只匹配了4个令牌中的2个。
时刻4：请求0完成并被释放。块2、3和4按逆序加入空闲队列（但块2和3仍处于缓存状态）。块0和1未被加入空闲队列，因为它们正被请求1使用。
时刻5：请求1完成并被释放。
时刻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（被淘汰）。