1. 项目概述为什么我们需要关注 LLM 推理引擎如果你最近在部署或使用大语言模型LLM大概率听过 vLLM 这个名字。它几乎成了高效 LLM 推理的代名词。但当你真正去阅读它的源码或尝试深度定制时可能会被其复杂的调度、内存管理和分布式逻辑所震撼。这时一个更轻量、更聚焦于核心原理的实现——比如Nano-vLLM——就成了绝佳的学习样本。它不是要替代 vLLM而是像一张清晰的解剖图帮你剥离繁杂的工程外壳直击 LLM 推理引擎最核心的几块“骨骼”调度器Scheduler、注意力Attention计算优化、以及 KV Cache 的内存管理。简单来说LLM 推理引擎的核心任务就一个用有限的硬件资源主要是 GPU 显存和算力以尽可能高的吞吐量Tokens per Second和尽可能低的延迟服务好用户的推理请求。这听起来像是一个经典的资源调度问题但 LLM 的自回归生成特性下一个 Token 依赖于之前所有 Token和巨大的模型参数让这个问题变得异常棘手。vLLM 提出的PagedAttention和其高效的调度器是解决这个问题的关键创新。而 Nano-vLLM 则试图用最精简的代码复现这一核心思想让我们能亲手“摸到”这些机制的运行脉络。这篇文章我们就以 Nano-vLLM 为透镜深入 LLM 推理引擎的内部。无论你是希望优化自家模型的部署效率还是单纯对底层技术充满好奇理解这些原理都将让你在设计和排查问题时拥有更清晰的视角。我们将从最根本的调度逻辑开始一步步拆解一个推理请求是如何被处理、计算并最终生成文本的。2. 核心架构与设计思路拆解一个完整的 LLM 推理引擎可以抽象为几个相互协作的组件。Nano-vLLM 为了教学清晰对其进行了高度简化但保留了最关键的链路。2.1 核心组件交互全景在一个简化的视图里推理请求的旅程是这样的用户请求入口用户发送一个包含提示词Prompt的请求。调度器Scheduler这是引擎的大脑。它接收请求决定何时、在哪个计算核心上执行它。它管理着所有等待中和运行中的请求队列并处理一个关键难题如何让多个请求共享 GPU 计算资源尤其是显存。模型执行器Model Executor这是引擎的肌肉。它接收调度器分配好的、一批待处理的请求调用底层模型如 Llama、Qwen 的 Transformer 块进行前向计算。这里涉及的关键优化是Attention 的批量计算和KV Cache 的管理。内存管理器Memory Manager这是引擎的仓库。它专门负责 KV Cache 这块“动态内存”的分配、释放和共享。vLLM 的革命性创新PagedAttention就是在这里实现的它让 KV Cache 可以像操作系统管理内存一样以“页”为单位进行灵活管理极大提高了显存利用率。Nano-vLLM 的设计思路是分而治之。它将调度、计算、内存管理解耦让我们可以单独研究每个部分。例如你可以先实现一个最简单的先进先出FIFO调度器然后再加入更复杂的、支持中断继续的调度策略。这种模块化设计正是学习复杂系统的最佳路径。2.2 为什么是“分页”式 KV Cache要理解调度器和内存管理器在忙什么必须首先理解 KV Cache 的重要性。在自回归生成中模型在计算第t个 token 时需要用到之前所有1到t-1个 token 对应的 Key 和 Value 向量即 K, V。这些向量如果每次都重新计算开销巨大。因此标准的做法是把它们缓存起来这就是 KV Cache。问题来了每个请求的生成序列长度是动态的、不可预知的。如果为每个请求预先分配一个可能的最大长度比如 2048对于短请求将是巨大的浪费如果分配不足长请求又会失败。此外多个请求之间的 Cache 无法共享即使它们的提示词前缀完全相同。注意这就是传统动态显存分配面临的“外部碎片”问题。频繁地分配和释放不同大小的内存块会在显存中留下许多无法被利用的小空隙。vLLM 的 PagedAttention 借鉴了操作系统的虚拟内存和分页思想将 KV Cache 空间划分为固定大小的“块”Block比如每个块存储 16 个 token 的 K 和 V。每个请求的 KV Cache 被视为由一系列这样的“块”组成的逻辑空间。内存管理器维护一个全局的空闲块列表。当一个请求需要更多空间来存储新生成的 token 的 KV 时就从空闲列表中分配一个或多个物理块给它。当请求结束时它占用的所有块被归还到空闲列表供其他请求使用。这样做的好处是显而易见的消除了外部碎片。因为所有分配单元大小相同任何空闲块都可以满足任何请求的分配需求。同时它为实现请求间的Memory Sharing奠定了基础——如果两个请求有相同的提示词前缀它们可以指向同一组物理块从而节省大量显存。Nano-vLLM 的核心目标之一就是用可读的代码展示这一分页机制是如何从数据结构层面建立起来的。3. 调度器Scheduler深度解析调度器是推理服务高吞吐、低延迟的指挥中枢。它的决策直接影响了 GPU 的利用率和用户的等待时间。3.1 调度器的基本职责与策略调度器持续监控两个队列等待队列Pending Queue和运行队列Running Queue。它的核心循环是检查是否有新请求到达放入等待队列。根据某种策略从等待队列中选择一个或多个请求将其移入运行队列并为其分配计算资源主要是 GPU 算力和 KV Cache 块。触发模型执行器对运行队列中的所有请求进行一步一个 Token的计算。处理计算完成后的请求如果请求生成结束则释放其所有资源否则等待下一轮调度。最简单的策略是First-Come-First-Served (FCFS)。但这对长短请求混合的场景不友好一个长请求会阻塞后面所有短请求造成“队头阻塞”。更先进的调度器如 vLLM 默认采用的是一种Continuous Batching策略。它允许迭代级调度每个解码步生成一个 Token都可以重新调度。请求的挂起与恢复如果一个运行中的请求在当前步暂时无法获得资源比如 KV Cache 块不足它可以被挂起让其他可以运行的请求先执行。细粒度资源管理调度决策基于当前可用的精确资源空闲块数、GPU 计算单元做出。Nano-vLLM 通常会实现一个简化版的连续批处理调度器。其关键数据结构可能包括一个为每个请求维护的“状态机”记录它当前解码到了哪一步、占用了哪些物理块、以及是否处于可运行状态。3.2 调度与内存管理的协同调度器不能独自做决定。在决定将哪些请求加入运行队列前它必须咨询内存管理器“如果我要运行这几个请求的下一个 Token我们需要多少新的 KV Cache 块当前空闲块够吗”这个过程称为预分配Pre-allocation或计划Planning。调度器会模拟一次调度决策向内存管理器申请所需的块。如果内存管理器批准即空闲块足够则调度生效如果不足调度器可能需要调整策略例如只选择部分请求运行或者挂起某些已运行但需要新块的请求。这种紧密的协同确保了系统永远不会在运行时因为显存不足而崩溃同时也实现了资源利用率的最大化。在 Nano-vLLM 的代码中你可能会看到一个Scheduler类持有一个MemoryManager的引用并在其schedule()方法中频繁调用memory_manager.can_allocate(requests)这样的接口。4. 注意力计算与 KV Cache 管理实战理解了调度逻辑我们再看模型实际是如何计算的。这部分是性能的关键涉及大量的 GPU 编程优化。4.1 PagedAttention 的前向计算实现传统的 Attention 计算要求 K 和 V 张量在内存中是连续的。但 PagedAttention 中一个请求的 KV Cache 可能分散在多个不连续的物理块中。因此我们需要一个特殊的 Attention 算子它能够根据一个“块表”来 gather 分散的 K 和 V。假设我们有一个请求它的逻辑 KV Cache 长度是L块大小是B。那么它需要ceil(L / B)个物理块。我们用一个列表block_ids [3, 7, 12, ...]来记录这些物理块的 ID。在计算第t个 token 的 Attention 时t小于L我们需要取出前t个 token 对应的 K 和 V。这些数据分布在block_ids指向的各个物理块中。我们需要计算第t个 token 落在哪个物理块block_idx t // B以及在该块内的偏移offset t % B。实际的 GPU 核函数会接收所有物理块组成的大张量、block_ids列表以及请求的序列信息通过一次高效的内存访问 gather 出这个请求所需的、连续的 K 和 V 张量再进行标准的 Attention 计算。Nano-vLLM 为了简化可能会先用一个 CPU 模拟版本实现这个 gather 逻辑让你理解其数据流。真正的 vLLM 则使用了高度优化的 CUDA 内核将 gather 和 Attention 计算融合以最小化内存带宽的消耗。4.2 内存管理器的数据结构与算法内存管理器是 PagedAttention 的基石。它的核心数据结构通常包括Block表示一个固定大小的物理内存块。包含一个唯一 ID 和存储的数据K V。FreeBlockPool一个空闲物理块的列表或堆。初始时所有块都在这里。AllocatedBlocks一个映射记录每个请求 ID 分配了哪些物理块。其关键操作很简单allocate(seq, num_blocks)为序列seq分配num_blocks个物理块。从空闲池取出记录分配关系返回块 ID 列表。free(seq)释放序列seq占用的所有物理块将其归还空闲池。can_allocate(num_blocks)查询当前是否有足够num_blocks个空闲块。为了实现块共享Prefix Caching还需要更复杂的数据结构比如一个基于内容哈希的块索引。当一个新的请求到来时先将其提示词的哈希值与已有块的哈希值对比如果匹配则直接让该请求的块表指向已有的物理块而不是分配新块。Nano-vLLM 可能会演示这个机制的基本原理。实操心得在实现内存管理器时锁的粒度是需要仔细考虑的问题。调度器和多个工作线程可能并发地申请和释放块。一个全局大锁会限制性能但过于细粒度的锁又容易引入死锁。一个常见的折中方案是为空闲池和每个请求的分配表使用不同的锁。5. 从零开始构建一个极简推理引擎现在让我们把理论付诸实践勾勒出构建一个类似 Nano-vLLM 的极简推理引擎的步骤。这能帮你把散落的知识点串联起来。5.1 第一步定义核心数据结构首先我们需要用代码定义出我们的“世界”。# 定义物理块。在实际中它对应GPU显存中的一块区域。 class PhysicalBlock: def __init__(self, block_id: int, block_size: int): self.block_id block_id self.k_data torch.zeros((block_size, hidden_size)) # 模拟K缓存 self.v_data torch.zeros((block_size, hidden_size)) # 模拟V缓存 self.ref_count 0 # 引用计数用于块共享 # 内存管理器 class MemoryManager: def __init__(self, total_blocks: int, block_size: int): self.free_blocks [PhysicalBlock(i, block_size) for i in range(total_blocks)] self.allocated {} # seq_id - list[PhysicalBlock] def allocate_for_seq(self, seq_id, num_blocks): if len(self.free_blocks) num_blocks: return None # 分配失败 allocated self.free_blocks[:num_blocks] self.free_blocks self.free_blocks[num_blocks:] self.allocated[seq_id] allocated return allocated def free_seq(self, seq_id): for block in self.allocated.get(seq_id, []): block.ref_count - 1 if block.ref_count 0: self.free_blocks.append(block) self.allocated.pop(seq_id, None) # 请求序列 class Sequence: def __init__(self, seq_id: int, prompt: str): self.seq_id seq_id self.prompt_ids encode(prompt) self.generated_ids [] self.block_table [] # 记录本序列使用的物理块ID列表 self.status WAITING # WAITING, RUNNING, FINISHED5.2 第二步实现调度循环接着我们实现一个单线程的、简化版的调度循环。class SimpleScheduler: def __init__(self, memory_manager: MemoryManager, max_running_seq: int): self.waiting_queue [] self.running_queue [] self.memory_manager memory_manager self.max_running_seq max_running_seq def add_request(self, seq: Sequence): self.waiting_queue.append(seq) def schedule_step(self): # 1. 尝试将等待队列的请求加入运行队列 while len(self.running_queue) self.max_running_seq and self.waiting_queue: seq self.waiting_queue.pop(0) # 预估该序列下一步需要多少新块例如生成第一个token需要为prompt分配块 needed_blocks estimate_blocks_needed(seq) allocated self.memory_manager.allocate_for_seq(seq.seq_id, needed_blocks) if allocated: seq.block_table [b.block_id for b in allocated] seq.status RUNNING self.running_queue.append(seq) else: # 内存不足放回等待队列头部 self.waiting_queue.insert(0, seq) break # 无法调度更多 # 2. 执行运行队列中所有序列的一步解码 if self.running_queue: # 这里会调用模型执行器进行批量前向计算 # 假设 model_step 会更新每个seq的generated_ids finished_seqs model_step(self.running_queue) # 3. 处理已完成的序列释放资源 for seq in finished_seqs: seq.status FINISHED self.memory_manager.free_seq(seq.seq_id) self.running_queue.remove(seq)这个循环虽然简单但已经包含了调度选择哪些请求运行、资源管理分配块、计算model_step和回收释放块的全流程。5.3 第三步集成注意力计算最后我们需要在model_step函数中实现支持分页的注意力计算。这里展示其核心逻辑的伪代码def paged_attention(query, # 当前token的查询向量 [batch, hidden] block_tables, # 每个序列的块表列表 k_cache, # 所有物理块的K缓存大张量 [total_blocks, block_size, hidden] v_cache, # 所有物理块的V缓存大张量 seq_lengths): # 每个序列当前的总长度prompt generated batch_size query.shape[0] scores [] outputs [] for i in range(batch_size): seq_len seq_lengths[i] block_table block_tables[i] # 1. Gather: 根据块表和序列长度从物理缓存中取出该序列所需的连续K, V # 计算需要哪些块以及块内偏移 k_seq gather_from_blocks(k_cache, block_table, seq_len) v_seq gather_from_blocks(v_cache, block_table, seq_len) # 2. 标准Attention计算 attn_scores torch.matmul(query[i].unsqueeze(0), k_seq.transpose(-1, -2)) attn_weights F.softmax(attn_scores, dim-1) out torch.matmul(attn_weights, v_seq).squeeze(0) outputs.append(out) return torch.stack(outputs)gather_from_blocks函数是这个过程的关键它实现了从非连续物理块到逻辑连续张量的转换。在真实的 GPU 实现中这一步会通过自定义内核高效完成。6. 常见问题与性能调优实战在理解和实现基础版本后你会遇到更实际的问题。以下是一些典型场景和排查思路。6.1 吞吐量上不去检查你的调度与计算瓶颈问题现象GPU 利用率低生成速度慢吞吐量远低于预期。排查思路调度器是否成为瓶颈在 CPU 上运行的调度逻辑如果过于复杂可能赶不上 GPU 的计算速度。可以尝试简化调度策略或者将调度器本身的部分工作如块表管理移到 GPU 上。批处理大小Batch Size是否过小GPU 擅长大规模并行计算。如果运行队列中始终只有一两个请求在计算GPU 的算力就被浪费了。可以尝试调整max_running_seq但要注意这会增加显存压力。注意力计算是瓶颈吗使用nsys或nvprof等性能分析工具查看 GPU 内核的执行时间。如果PagedAttention的自定义内核耗时很长可能需要检查其实现效率或者考虑是否因块过于分散导致内存访问效率低下缓存命中率低。调优技巧实现一个流水线Pipeline。将调度、数据准备将输入 token 从 CPU 搬到 GPU、模型计算、结果回写将生成的 token 从 GPU 搬回 CPU这几个阶段重叠起来。当 GPU 在执行第 N 批的计算时CPU 已经在为第 N1 批准备数据了。这能有效隐藏数据搬运的开销。6.2 显存溢出OOM分析你的内存管理问题现象在运行一段时间后或处理特定长序列时出现 CUDA out of memory 错误。排查思路内存泄漏这是最常见的原因。确保每个请求结束后其占用的所有物理块都被正确释放。在MemoryManager.free_seq方法中加入详细的日志跟踪块的分配和释放是否成对出现。碎片化问题即使使用分页如果块大小设置不当也可能造成内部碎片。例如块大小为 16但大量请求的序列长度都是 17那么每个请求都需要 2 个块第二个块只用了 1 个位置浪费了 15 个位置。可以尝试分析请求的长度分布调整块大小。共享失效预期的前缀共享没有发生。检查哈希函数是否合理以及共享逻辑是否正确。一个常见错误是只对完整的提示词进行哈希而忽略了中间生成结果的共享可能性。调优技巧实现一个块的重用策略。当空闲块池耗尽时不要立即失败可以尝试将一些暂时不活跃例如被调度器挂起的请求的块交换到 CPU 内存如果系统内存足够大腾出 GPU 显存。这类似于操作系统的“交换Swap”是一种用时间换空间的策略。6.3 生成结果不一致或错误调试你的计算逻辑问题现象生成的文本不符合预期或者与标准 Transformers 库的输出不一致。排查思路Gather 逻辑错误这是最可能出问题的地方。写一个单元测试构造一个简单的、已知的 KV Cache 分布场景手动计算期望的 Attention 输出与你的paged_attention函数结果对比。重点检查块表索引和块内偏移的计算。状态管理混乱确保每个序列的block_table和当前生成位置position是严格同步的。在调度器挂起和恢复一个序列时这些状态必须被完美保存和恢复。数值精度问题在 GPU 上混合使用 FP16 和 FP32 可能导致细微的精度差异经过多步生成后放大。确保你的模型权重、输入数据和缓存数据精度一致。调试技巧实现一个“验证模式”。在关键步骤如调度决策后、注意力计算前将张量数据 dump 下来与一个已知正确的参考实现如 Hugging Face 的transformers库以非优化模式运行的中间结果进行逐元素对比。这能帮你快速定位首次出现偏差的环节。理解 Nano-vLLM 或类似教学项目的意义不在于复制一个生产级的推理引擎而在于亲手搭建起核心组件的骨架感受数据在调度器、内存管理器和计算内核间的流动。当你再去看 vLLM、TGI 这些成熟项目的源码时那些复杂的工程细节就不再是黑盒而是你已理解的骨架之上为了极致性能、鲁棒性和功能丰富性而添加的血肉。这正是深入理解 LLM 推理引擎的第一步也是最坚实的一步。