vLLM推理引擎架构PagedAttention机制与高吞吐推理一、大模型推理的显存碎片困境KV Cache的管理挑战大模型推理的核心瓶颈在于KV Cache的显存管理。自回归生成过程中模型需要缓存每一步的Key和Value向量用于后续Token的注意力计算。KV Cache的大小与序列长度和批量大小成正比——一个70B模型处理4096长度的序列时单个请求的KV Cache可达数GB。传统的推理引擎为每个请求预分配一块连续的显存来存储KV Cache这导致严重的显存碎片问题预分配过大会浪费显存预分配过小则序列生成到一半时需要重新分配引发延迟尖峰不同请求的序列长度差异大预分配大小难以统一请求完成后释放的显存块可能无法被新请求复用外部碎片。vLLM通过PagedAttention机制借鉴操作系统的虚拟内存分页思想将KV Cache按固定大小的Block管理按需分配和释放彻底解决了显存碎片问题。二、PagedAttention机制详解2.1 核心思想KV Cache的分页管理PagedAttention将每个请求的KV Cache组织为多个固定大小的Block每个Block存储固定数量Token的Key和Value向量。Block按需分配不需要连续的显存空间类似于操作系统的分页内存管理。graph TB subgraph 传统预分配 A1[请求1: 连续显存块br/大量浪费空间] A2[请求2: 连续显存块br/空间不足需扩容] A3[请求3: 等待分配br/碎片无法利用] end subgraph PagedAttention B1[请求1: Block列表br/按需分配] B2[请求2: Block列表br/按需分配] B3[请求3: Block列表br/复用已释放Block] end subgraph Block池 C[Block 0][Block 1][Block 2][Block 3] D[Block 4][Block 5][Block 6][Block 7] end B1 -- C B1 -- D B2 -- C B3 -- C2.2 Block管理器实现class BlockManager: KV Cache Block管理器 def __init__(self, num_blocks: int, block_size: int, num_heads: int, head_dim: int, dtypetorch.float16): self.block_size block_size # 每个Block存储的Token数 self.num_blocks num_blocks self.free_blocks list(range(num_blocks)) # 空闲Block列表 # 预分配KV Cache的Block池 element_size torch.tensor([], dtypedtype).element_size() self.k_cache torch.zeros( num_blocks, block_size, num_heads, head_dim, dtypedtype, devicecuda ) self.v_cache torch.zeros( num_blocks, block_size, num_heads, head_dim, dtypedtype, devicecuda ) def allocate(self, num_blocks: int) - list: 分配指定数量的Block if len(self.free_blocks) num_blocks: raise OutOfMemoryError( fNeed {num_blocks} blocks, fonly {len(self.free_blocks)} available ) allocated self.free_blocks[:num_blocks] self.free_blocks self.free_blocks[num_blocks:] return allocated def free(self, block_ids: list): 释放Block回空闲池 self.free_blocks.extend(block_ids) def get_block_table(self, request_id: str) - list: 获取请求的Block映射表 return self.request_block_map.get(request_id, [])2.3 PagedAttention KernelPagedAttention的核心挑战是注意力计算需要访问KV Cache但KV Cache不再存储在连续的显存地址中而是分散在多个Block中。需要自定义CUDA Kernel实现跨Block的注意力计算。class PagedAttentionFunction(torch.autograd.Function): PagedAttention前向传播 staticmethod def forward(ctx, query, key_cache, value_cache, block_tables, context_lens, block_size): Args: query: [num_tokens, num_heads, head_dim] key_cache: [num_blocks, block_size, num_heads, head_dim] value_cache: [num_blocks, block_size, num_heads, head_dim] block_tables: [num_seqs, max_num_blocks_per_seq] context_lens: [num_seqs] 每个序列的实际长度 block_size: 每个Block的Token数 num_tokens query.shape[0] num_heads query.shape[1] head_dim query.shape[2] output torch.empty_like(query) # 调用自定义CUDA Kernel # 核心逻辑根据block_tables找到每个Token对应的KV Block # 在Block内计算注意力分数跨Block累加 paged_attention_kernel( output, query, key_cache, value_cache, block_tables, context_lens, block_size ) return output三、连续批处理与调度3.1 连续批处理Continuous Batching传统推理引擎采用静态批处理——等待所有请求完成后才开始下一批。vLLM采用连续批处理Iteration-level Scheduling每次迭代都重新调度已完成的请求移出新请求加入无需等待同批其他请求完成。class Scheduler: 请求调度器 def __init__(self, block_manager: BlockManager, max_num_seqs: int 256): self.block_manager block_manager self.max_num_seqs max_num_seqs self.waiting_queue [] # 等待队列 self.running_seqs [] # 正在运行的序列 self.max_seq_len 8192 def schedule(self) - SchedulerOutput: 每次迭代调度决定哪些序列参与本轮计算 scheduled_seqs [] # 1. 保留正在运行的序列 for seq in self.running_seqs: if seq.is_finished(): # 序列已完成释放Block self.block_manager.free(seq.block_table) else: # 检查是否需要分配新的Block if self._need_new_block(seq): new_block self.block_manager.allocate(1) if new_block: seq.block_table.extend(new_block) scheduled_seqs.append(seq) else: # 显存不足抢占Preemption self._preempt(seq) else: scheduled_seqs.append(seq) # 2. 从等待队列中添加新序列 remaining_slots self.max_num_seqs - len(scheduled_seqs) for _ in range(remaining_slots): if not self.waiting_queue: break seq self.waiting_queue.pop(0) # 为新序列分配初始Block num_blocks math.ceil( seq.prompt_len / self.block_manager.block_size) blocks self.block_manager.allocate(num_blocks) if blocks: seq.block_table blocks scheduled_seqs.append(seq) else: # 显存不足放回等待队列 self.waiting_queue.insert(0, seq) break self.running_seqs scheduled_seqs return SchedulerOutput(running_seqsscheduled_seqs) def _preempt(self, seq): 抢占策略释放低优先级序列的Block # 通过重新计算Recomputation而非交换Swapping恢复 self.block_manager.free(seq.block_table) seq.block_table [] seq.num_computed_tokens 0 # 需要重新计算 self.waiting_queue.append(seq)3.2 抢占与恢复策略当显存不足时调度器需要抢占Preempt部分运行中的序列释放其Block给更高优先级的序列。vLLM采用重新计算策略——被抢占的序列放回等待队列重新执行Prompt阶段恢复KV Cache。四、架构权衡与边界分析4.1 Block大小的选择Block过小如1个TokenBlock表的管理开销增大注意力Kernel的访存效率降低Block过大如256个Token预分配浪费增加短序列的显存利用率下降。vLLM默认Block大小为16在管理开销和利用率之间取得平衡。4.2 抢占策略的开销重新计算策略在抢占频繁时会浪费大量计算资源。当系统过载时抢占-恢复的循环可能导致整体吞吐量下降。建议设置合理的最大并发序列数避免过度调度。4.3 前缀缓存的复用多个请求可能共享相同的System Prompt前缀PagedAttention支持将共享前缀的KV Cache Block标记为只读多个请求复用同一组Block显著降低显存占用和重复计算。但前缀匹配的检测逻辑增加了调度器的复杂度。五、总结vLLM通过PagedAttention机制将KV Cache按Block管理按需分配和释放解决了传统推理引擎的显存碎片问题。连续批处理实现了迭代级调度提升了GPU利用率。抢占与恢复策略在显存紧张时保障系统稳定性。落地建议从默认Block大小16开始根据实际序列长度分布调整监控抢占频率高频抢占时降低最大并发数启用前缀缓存复用共享System Prompt但需评估匹配检测的额外开销。
vLLM推理引擎架构：PagedAttention机制与高吞吐推理
vLLM推理引擎架构PagedAttention机制与高吞吐推理一、大模型推理的显存碎片困境KV Cache的管理挑战大模型推理的核心瓶颈在于KV Cache的显存管理。自回归生成过程中模型需要缓存每一步的Key和Value向量用于后续Token的注意力计算。KV Cache的大小与序列长度和批量大小成正比——一个70B模型处理4096长度的序列时单个请求的KV Cache可达数GB。传统的推理引擎为每个请求预分配一块连续的显存来存储KV Cache这导致严重的显存碎片问题预分配过大会浪费显存预分配过小则序列生成到一半时需要重新分配引发延迟尖峰不同请求的序列长度差异大预分配大小难以统一请求完成后释放的显存块可能无法被新请求复用外部碎片。vLLM通过PagedAttention机制借鉴操作系统的虚拟内存分页思想将KV Cache按固定大小的Block管理按需分配和释放彻底解决了显存碎片问题。二、PagedAttention机制详解2.1 核心思想KV Cache的分页管理PagedAttention将每个请求的KV Cache组织为多个固定大小的Block每个Block存储固定数量Token的Key和Value向量。Block按需分配不需要连续的显存空间类似于操作系统的分页内存管理。graph TB subgraph 传统预分配 A1[请求1: 连续显存块br/大量浪费空间] A2[请求2: 连续显存块br/空间不足需扩容] A3[请求3: 等待分配br/碎片无法利用] end subgraph PagedAttention B1[请求1: Block列表br/按需分配] B2[请求2: Block列表br/按需分配] B3[请求3: Block列表br/复用已释放Block] end subgraph Block池 C[Block 0][Block 1][Block 2][Block 3] D[Block 4][Block 5][Block 6][Block 7] end B1 -- C B1 -- D B2 -- C B3 -- C2.2 Block管理器实现class BlockManager: KV Cache Block管理器 def __init__(self, num_blocks: int, block_size: int, num_heads: int, head_dim: int, dtypetorch.float16): self.block_size block_size # 每个Block存储的Token数 self.num_blocks num_blocks self.free_blocks list(range(num_blocks)) # 空闲Block列表 # 预分配KV Cache的Block池 element_size torch.tensor([], dtypedtype).element_size() self.k_cache torch.zeros( num_blocks, block_size, num_heads, head_dim, dtypedtype, devicecuda ) self.v_cache torch.zeros( num_blocks, block_size, num_heads, head_dim, dtypedtype, devicecuda ) def allocate(self, num_blocks: int) - list: 分配指定数量的Block if len(self.free_blocks) num_blocks: raise OutOfMemoryError( fNeed {num_blocks} blocks, fonly {len(self.free_blocks)} available ) allocated self.free_blocks[:num_blocks] self.free_blocks self.free_blocks[num_blocks:] return allocated def free(self, block_ids: list): 释放Block回空闲池 self.free_blocks.extend(block_ids) def get_block_table(self, request_id: str) - list: 获取请求的Block映射表 return self.request_block_map.get(request_id, [])2.3 PagedAttention KernelPagedAttention的核心挑战是注意力计算需要访问KV Cache但KV Cache不再存储在连续的显存地址中而是分散在多个Block中。需要自定义CUDA Kernel实现跨Block的注意力计算。class PagedAttentionFunction(torch.autograd.Function): PagedAttention前向传播 staticmethod def forward(ctx, query, key_cache, value_cache, block_tables, context_lens, block_size): Args: query: [num_tokens, num_heads, head_dim] key_cache: [num_blocks, block_size, num_heads, head_dim] value_cache: [num_blocks, block_size, num_heads, head_dim] block_tables: [num_seqs, max_num_blocks_per_seq] context_lens: [num_seqs] 每个序列的实际长度 block_size: 每个Block的Token数 num_tokens query.shape[0] num_heads query.shape[1] head_dim query.shape[2] output torch.empty_like(query) # 调用自定义CUDA Kernel # 核心逻辑根据block_tables找到每个Token对应的KV Block # 在Block内计算注意力分数跨Block累加 paged_attention_kernel( output, query, key_cache, value_cache, block_tables, context_lens, block_size ) return output三、连续批处理与调度3.1 连续批处理Continuous Batching传统推理引擎采用静态批处理——等待所有请求完成后才开始下一批。vLLM采用连续批处理Iteration-level Scheduling每次迭代都重新调度已完成的请求移出新请求加入无需等待同批其他请求完成。class Scheduler: 请求调度器 def __init__(self, block_manager: BlockManager, max_num_seqs: int 256): self.block_manager block_manager self.max_num_seqs max_num_seqs self.waiting_queue [] # 等待队列 self.running_seqs [] # 正在运行的序列 self.max_seq_len 8192 def schedule(self) - SchedulerOutput: 每次迭代调度决定哪些序列参与本轮计算 scheduled_seqs [] # 1. 保留正在运行的序列 for seq in self.running_seqs: if seq.is_finished(): # 序列已完成释放Block self.block_manager.free(seq.block_table) else: # 检查是否需要分配新的Block if self._need_new_block(seq): new_block self.block_manager.allocate(1) if new_block: seq.block_table.extend(new_block) scheduled_seqs.append(seq) else: # 显存不足抢占Preemption self._preempt(seq) else: scheduled_seqs.append(seq) # 2. 从等待队列中添加新序列 remaining_slots self.max_num_seqs - len(scheduled_seqs) for _ in range(remaining_slots): if not self.waiting_queue: break seq self.waiting_queue.pop(0) # 为新序列分配初始Block num_blocks math.ceil( seq.prompt_len / self.block_manager.block_size) blocks self.block_manager.allocate(num_blocks) if blocks: seq.block_table blocks scheduled_seqs.append(seq) else: # 显存不足放回等待队列 self.waiting_queue.insert(0, seq) break self.running_seqs scheduled_seqs return SchedulerOutput(running_seqsscheduled_seqs) def _preempt(self, seq): 抢占策略释放低优先级序列的Block # 通过重新计算Recomputation而非交换Swapping恢复 self.block_manager.free(seq.block_table) seq.block_table [] seq.num_computed_tokens 0 # 需要重新计算 self.waiting_queue.append(seq)3.2 抢占与恢复策略当显存不足时调度器需要抢占Preempt部分运行中的序列释放其Block给更高优先级的序列。vLLM采用重新计算策略——被抢占的序列放回等待队列重新执行Prompt阶段恢复KV Cache。四、架构权衡与边界分析4.1 Block大小的选择Block过小如1个TokenBlock表的管理开销增大注意力Kernel的访存效率降低Block过大如256个Token预分配浪费增加短序列的显存利用率下降。vLLM默认Block大小为16在管理开销和利用率之间取得平衡。4.2 抢占策略的开销重新计算策略在抢占频繁时会浪费大量计算资源。当系统过载时抢占-恢复的循环可能导致整体吞吐量下降。建议设置合理的最大并发序列数避免过度调度。4.3 前缀缓存的复用多个请求可能共享相同的System Prompt前缀PagedAttention支持将共享前缀的KV Cache Block标记为只读多个请求复用同一组Block显著降低显存占用和重复计算。但前缀匹配的检测逻辑增加了调度器的复杂度。五、总结vLLM通过PagedAttention机制将KV Cache按Block管理按需分配和释放解决了传统推理引擎的显存碎片问题。连续批处理实现了迭代级调度提升了GPU利用率。抢占与恢复策略在显存紧张时保障系统稳定性。落地建议从默认Block大小16开始根据实际序列长度分布调整监控抢占频率高频抢占时降低最大并发数启用前缀缓存复用共享System Prompt但需评估匹配检测的额外开销。
