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 --> C

2.2 Block管理器实现

class BlockManager: """KV Cache Block管理器""" def __init__(self, num_blocks: int, block_size: int, num_heads: int, head_dim: int, dtype=torch.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([], dtype=dtype).element_size() self.k_cache = torch.zeros( num_blocks, block_size, num_heads, head_dim, dtype=dtype, device='cuda' ) self.v_cache = torch.zeros( num_blocks, block_size, num_heads, head_dim, dtype=dtype, device='cuda' ) def allocate(self, num_blocks: int) -> list: """分配指定数量的Block""" if len(self.free_blocks) < num_blocks: raise OutOfMemoryError( f"Need {num_blocks} blocks, " f"only {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 Kernel

PagedAttention的核心挑战是：注意力计算需要访问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_seqs=scheduled_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)