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 核心组件交互全景

在一个简化的视图里，推理请求的旅程是这样的：

1. 用户请求入口：用户发送一个包含提示词（Prompt）的请求。
2. 调度器（Scheduler）：这是引擎的大脑。它接收请求，决定何时、在哪个计算核心上执行它。它管理着所有等待中和运行中的请求队列，并处理一个关键难题：如何让多个请求共享 GPU 计算资源，尤其是显存。
3. 模型执行器（Model Executor）：这是引擎的肌肉。它接收调度器分配好的、一批待处理的请求，调用底层模型（如 Llama、Qwen 的 Transformer 块）进行前向计算。这里涉及的关键优化是Attention 的批量计算和KV Cache 的管理。
4. 内存管理器（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）。它的核心循环是：

1. 检查是否有新请求到达，放入等待队列。
2. 根据某种策略，从等待队列中选择一个或多个请求，将其移入运行队列，并为其分配计算资源（主要是 GPU 算力和 KV Cache 块）。
3. 触发模型执行器对运行队列中的所有请求进行一步（一个 Token）的计算。
4. 处理计算完成后的请求：如果请求生成结束，则释放其所有资源；否则，等待下一轮调度。

最简单的策略是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）：

1. 我们需要取出前t个 token 对应的 K 和 V。
2. 这些数据分布在block_ids指向的各个物理块中。我们需要计算第t个 token 落在哪个物理块（block_idx = t // B），以及在该块内的偏移（offset = t % B）。
3. 实际的 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, FINISHED

5.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 利用率低，生成速度慢，吞吐量远低于预期。排查思路：

1. 调度器是否成为瓶颈？在 CPU 上运行的调度逻辑如果过于复杂，可能赶不上 GPU 的计算速度。可以尝试简化调度策略，或者将调度器本身的部分工作（如块表管理）移到 GPU 上。
2. 批处理大小（Batch Size）是否过小？GPU 擅长大规模并行计算。如果运行队列中始终只有一两个请求在计算，GPU 的算力就被浪费了。可以尝试调整max_running_seq，但要注意这会增加显存压力。
3. 注意力计算是瓶颈吗？使用nsys或nvprof等性能分析工具，查看 GPU 内核的执行时间。如果PagedAttention的自定义内核耗时很长，可能需要检查其实现效率，或者考虑是否因块过于分散导致内存访问效率低下（缓存命中率低）。

调优技巧：实现一个流水线（Pipeline）。将调度、数据准备（将输入 token 从 CPU 搬到 GPU）、模型计算、结果回写（将生成的 token 从 GPU 搬回 CPU）这几个阶段重叠起来。当 GPU 在执行第 N 批的计算时，CPU 已经在为第 N+1 批准备数据了。这能有效隐藏数据搬运的开销。

6.2 显存溢出（OOM）？分析你的内存管理

问题现象：在运行一段时间后，或处理特定长序列时，出现 CUDA out of memory 错误。排查思路：

1. 内存泄漏：这是最常见的原因。确保每个请求结束后，其占用的所有物理块都被正确释放。在MemoryManager.free_seq方法中加入详细的日志，跟踪块的分配和释放是否成对出现。
2. 碎片化问题：即使使用分页，如果块大小设置不当，也可能造成内部碎片。例如，块大小为 16，但大量请求的序列长度都是 17，那么每个请求都需要 2 个块，第二个块只用了 1 个位置，浪费了 15 个位置。可以尝试分析请求的长度分布，调整块大小。
3. 共享失效：预期的前缀共享没有发生。检查哈希函数是否合理，以及共享逻辑是否正确。一个常见错误是只对完整的提示词进行哈希，而忽略了中间生成结果的共享可能性。

调优技巧：实现一个块的重用策略。当空闲块池耗尽时，不要立即失败，可以尝试将一些暂时不活跃（例如，被调度器挂起）的请求的块交换到 CPU 内存（如果系统内存足够大），腾出 GPU 显存。这类似于操作系统的“交换（Swap）”，是一种用时间换空间的策略。

6.3 生成结果不一致或错误？调试你的计算逻辑

问题现象：生成的文本不符合预期，或者与标准 Transformers 库的输出不一致。排查思路：

1. Gather 逻辑错误：这是最可能出问题的地方。写一个单元测试，构造一个简单的、已知的 KV Cache 分布场景，手动计算期望的 Attention 输出，与你的paged_attention函数结果对比。重点检查块表索引和块内偏移的计算。
2. 状态管理混乱：确保每个序列的block_table和当前生成位置position是严格同步的。在调度器挂起和恢复一个序列时，这些状态必须被完美保存和恢复。
3. 数值精度问题：在 GPU 上，混合使用 FP16 和 FP32 可能导致细微的精度差异，经过多步生成后放大。确保你的模型权重、输入数据和缓存数据精度一致。

调试技巧：实现一个“验证模式”。在关键步骤（如调度决策后、注意力计算前）将张量数据 dump 下来，与一个已知正确的参考实现（如 Hugging Face 的transformers库，以非优化模式运行）的中间结果进行逐元素对比。这能帮你快速定位首次出现偏差的环节。

理解 Nano-vLLM 或类似教学项目的意义，不在于复制一个生产级的推理引擎，而在于亲手搭建起核心组件的骨架，感受数据在调度器、内存管理器和计算内核间的流动。当你再去看 vLLM、TGI 这些成熟项目的源码时，那些复杂的工程细节就不再是黑盒，而是你已理解的骨架之上，为了极致性能、鲁棒性和功能丰富性而添加的血肉。这，正是深入理解 LLM 推理引擎的第一步，也是最坚实的一步。