深入解析LLM推理引擎：从PagedAttention到调度优化

1. 从“推理”到“引擎”为什么我们需要专门的LLM推理工具如果你最近在折腾大语言模型尤其是尝试自己部署一个7B、13B甚至更大参数的模型来跑点东西大概率会遇到一个场景你兴冲冲地下载了模型权重用上了最熟悉的PyTorch或者Transformers库写了几行加载和推理的代码。第一次生成回答时你满怀期待然后眼睁睁看着进度条缓慢蠕动GPU的显存占用瞬间飙升而生成几十个token所花的时间足够你泡一杯咖啡。这时你可能会想那些宣称“高性能”、“低延迟”的在线服务到底是怎么做到的秘密很大程度上就藏在“推理引擎”这四个字里。LLM推理远不止是model.generate()那么简单。它本质上是一个复杂的资源调度与计算优化问题。想象一下一个拥有数百亿参数的模型它的权重就像一座巨大的图书馆。每次推理生成一个token并不是要翻阅整个图书馆而是要根据当前的“阅读进度”已生成的文本序列快速定位到相关的几个书架激活的神经元或注意力头取出几本书执行矩阵运算然后写下下一句话。当有多个用户同时提问批处理或者一个用户要求生成长篇大论长序列时如何高效地管理这座“图书馆”的访客、减少重复的“找书”动作、避免“交通堵塞”就是推理引擎要解决的核心问题。原始的、通用的深度学习框架如PyTorch在设计上是为了灵活性和研发友好它提供了构建这座“图书馆”的所有砖瓦和图纸但并没有为“高峰期接待大量游客”这种特定场景做极致优化。于是像vLLM、TensorRT-LLM、TGIText Generation Inference等专门的LLM推理引擎应运而生。它们的目标非常明确在相同的硬件主要是GPU上用更少的显存生成更快的token同时服务更多的用户。而Nano-vLLM可以看作是vLLM的一个极简、可读的实现它剥离了生产级系统复杂的工程封装直指几个最核心的优化技术是理解这套“武功心法”的绝佳教材。所以这篇文章不是一份安装指南或API文档。我们将以Nano-vLLM为线索深入引擎盖之下看看现代LLM推理引擎是如何思考的。我们会聊到那个革命性的“显存救星”——PagedAttention会拆解调度器如何像交警一样指挥数据流动也会对比vLLM与其他方案的区别。无论你是算法工程师希望优化自家服务还是开发者好奇技术内幕抑或是学生想深入理解系统方向这篇“解剖报告”都将为你呈现LLM高效推理的核心逻辑与实现细节。2. 推理引擎的核心挑战显存、调度与计算在深入Nano-vLLM或vLLM的具体实现之前我们必须先建立共识一个高效的LLM推理引擎究竟在和什么作斗争理解了这些根本性的挑战我们才能明白后续每一个优化设计背后的深远用意。2.1 显存墙模型权重与KV Cache的双重压力LLM推理对显存的需求主要来自两部分模型参数Weights和KV Cache。模型参数是静态的在加载后基本不变。一个常见的衡量标准是FP16精度下每10亿参数大约需要2GB显存。一个70亿参数的模型仅权重就需要约14GB显存这已经接近甚至超过了许多消费级显卡如RTX 4090的24GB的容量上限更不用说更大的千亿级模型。KV Cache则是动态的、随着生成过程不断膨胀的“内存杀手”。在Transformer的解码生成阶段为了计算下一个token自注意力机制需要用到当前序列中所有之前token的Key和Value向量。这些向量被缓存起来避免每一轮生成都重新计算这就是KV Cache。它的可怕之处在于其增长与批次大小batch size和序列长度sequence length成正比。我们来算一笔账假设模型隐藏层维度为hidden_size4096注意力头数num_heads32那么每个注意力头的维度head_dim 4096/32128。对于序列中的一个token它在每一层都需要缓存一个Key向量和一个Value向量每个向量的大小是[num_heads, head_dim]。在FP16精度下一个token在一层中占用的KV Cache大小约为2 * 32 * 128 * 2 bytes ≈ 16 KB。对于一个有32层的模型一个token的KV Cache就膨胀到约16KB * 32 ≈ 512 KB。现在考虑一个中等规模的场景批次大小batch_size4生成序列长度seq_len512。那么总的KV Cache显存占用将达到4 * 512 * 512 KB ≈ 1 GB。这还只是缓存如果序列长度达到2048这个数字会变成4GB。在实时对话或长文档生成场景下KV Cache很容易成为显存使用的最大头甚至超过模型权重本身。传统的缓存管理方式是连续分配一整块内存。这带来了两个致命问题1内部碎片由于不同请求的序列长度不同为每个请求预留最大可能长度的内存会造成严重浪费2外部碎片当一些请求结束释放内存后留下的空闲内存块可能很小无法满足新请求的大内存需求导致虽然总空闲显存足够却无法分配。2.2 调度难题如何让GPU“忙”起来GPU是强大的并行计算设备但其计算能力能否被充分利用高度依赖于喂给它的数据是否“喂得饱”。LLM推理尤其是自回归生成存在天然的串行依赖必须算出第N个token才能基于它去计算第N1个token。这导致了严重的计算访存比低下问题每一轮生成GPU都要从显存中读取巨大的模型权重和KV Cache但实际执行的矩阵乘加运算量相对有限大部分时间花在了等待数据从显存搬运到计算单元上。此外在线服务中请求是随机到达的。有的用户问“你好”可能只需要生成10个token有的用户要求“写一篇小说”可能需要生成2000个token。如果来一个请求就立刻开始计算GPU很快就会因为处理短序列而处于“饥饿”状态计算单元闲置。如果等攒够一批请求再一起计算那么先来的用户就会经历漫长的等待延迟。因此调度器的核心使命是在吞吐量Throughput和延迟Latency之间寻找最佳平衡点。它需要决定何时将哪些请求的哪些token组合成一个批次Batch送入GPU计算如何管理这些请求的生命周期创建、挂起、恢复、完成当显存不足时如何做出取舍例如是否将部分请求的KV Cache交换到CPU内存。2.3 计算优化超越朴素的矩阵乘法即使解决了显存和调度问题计算本身也有巨大的优化空间。Transformer层中的注意力计算、前馈网络FFN计算都可以通过算子融合Kernel Fusion来减少内存读写开销。例如将LayerNorm、线性投影、激活函数等连续操作合并成一个CUDA Kernel可以避免中间结果写回显存再读出的过程。另外量化技术Quantization可以将模型权重和激活值从FP16降低到INT8甚至INT4从而显著减少显存占用和内存带宽压力虽然会引入一定的精度损失但在很多场景下是可以接受的权衡。现代推理引擎通常集成或提供了对接量化模型的通路。总结来说一个优秀的LLM推理引擎必须是一个“多面手”它既是显存管理大师能像操作系统管理内存一样精细地管理KV Cache又是调度策略专家能智能地编排请求最大化GPU利用率还是计算优化高手能榨干GPU的每一份算力。接下来我们就看看vLLM及其“教学版”Nano-vLLM是如何应对这些挑战的。3. PagedAttentionvLLM的“王牌”与灵感来源如果说vLLM只有一个技术点最值得称道那一定是PagedAttention。这个灵感来源于操作系统虚拟内存分页机制的思想从根本上重塑了KV Cache的管理方式解决了上一节提到的显存碎片化问题。理解PagedAttention是理解vLLM乃至所有现代高效推理引擎的关键。3.1 从操作系统虚拟内存获得的灵感在操作系统中每个进程都认为自己拥有一大片连续的物理内存。但实际上物理内存被划分成固定大小的“页”Page如4KB进程的虚拟地址空间也被分成同样大小的“页”。通过页表Page Table进行映射操作系统可以将进程不连续的虚拟页映射到物理内存中可能也不连续的物理页上。这样做的好处是消除外部碎片物理内存只需按页分配无需寻找一大块连续空间。共享内存不同的进程可以将自己的虚拟页映射到同一块物理页实现数据共享。按需加载只有进程实际访问到的页才会被调入物理内存。PagedAttention将这一套机制完美地移植到了KV Cache的管理上。它将每个请求的KV Cache在逻辑上视为一个连续的“逻辑块”但在物理上将其分割成固定大小的块Block进行分配。每个块可以存储固定数量token的Key和Value向量例如一个块存16个token的数据。3.2 PagedAttention的工作机制让我们通过一个具体例子来看PagedAttention如何工作。假设块大小Block Size为4个token。请求到来一个请求到来需要生成序列。系统为其创建一个逻辑上的“KV Cache空间”。按需分配物理块当该请求生成第一个token时系统分配一个物理块比如块ID 0来存储这个token的KV数据。此时逻辑位置0对应物理块0。持续生成与分配随着生成继续第2、3、4个token的KV数据依次填入物理块0。当要生成第5个token时块0已满系统分配一个新的物理块块ID 1。逻辑位置4-7就映射到了物理块1。管理映射关系系统维护一个块表Block Table类似于操作系统的页表记录每个请求的逻辑块到物理块的映射关系。例如请求A的块表可能是[0, 1, 3, ...]表示它的第0个逻辑块对应物理块0第1个逻辑块对应物理块1第2个逻辑块对应物理块3物理块2可能分配给了其他请求。这种机制带来了立竿见影的好处高效的内存利用物理块是固定大小的分配和释放都非常快速完全消除了外部碎片。内部碎片仅限于最后一个未满的块浪费极小。内存共享的基石这是PagedAttention更精妙的一环。在LLM推理中多个请求可能具有相同的前缀例如相同的系统提示词。在传统方式下每个请求都需要独立存储这份前缀的KV Cache。而在PagedAttention下这些请求可以共享存储相同前缀内容的物理块。系统只需在它们的块表中将对应的逻辑块映射到同一个物理块即可。这为多用户、多会话场景节省了大量显存。并行计算的优化由于块是固定大小的在组织注意力计算时可以将不同请求中“对齐”的块一起处理使得GPU核函数Kernel的编写更规整并行效率更高。3.3 在Nano-vLLM中窥见PagedAttention的实现Nano-vLLM的代码清晰地展示了这一核心思想。虽然为了简洁和可读性做了简化但骨架完整。你通常会看到类似如下的数据结构class Block: def __init__(self, block_id: int, block_size: int): self.block_id block_id self.block_size block_size self.ref_count 0 # 引用计数用于共享管理 self.token_ids [] # 存储这个块中的token id # 实际存储K、V张量的内存区域这里用列表示意 self.k_data [None] * block_size self.v_data [None] * block_size class BlockManager: def __init__(self, total_blocks: int, block_size: int): self.free_blocks list(range(total_blocks)) # 空闲块列表 self.allocated_blocks {} # block_id - Block 对象 self.block_size block_size def allocate_block(self) - Block: 分配一个空闲的物理块 if not self.free_blocks: raise OutOfMemoryError(No free blocks available) block_id self.free_blocks.pop() block Block(block_id, self.block_size) self.allocated_blocks[block_id] block return block def free_block(self, block: Block): 释放一个物理块当引用计数为0时 if block.ref_count 0: del self.allocated_blocks[block.block_id] self.free_blocks.append(block.block_id) class Request: def __init__(self, request_id: str): self.request_id request_id self.block_table [] # 记录逻辑块到物理Block对象的映射 self.current_position 0 # 当前生成到的逻辑位置在注意力计算时不再是直接操作一个大的、连续的KV Cache张量而是需要根据每个请求的block_table和current_position动态地收集Gather所有需要的物理块中的数据拼凑成这次计算所需的Key和Value张量。这个过程虽然引入了一些索引计算的开销但相比于它带来的显存利用率的巨大提升和碎片问题的解决这点开销是微不足道的。注意在实际的vLLM中块的管理、共享逻辑以及与之配套的注意力核函数如PagedAttention Kernel要复杂得多涉及高效的GPU原子操作和精细的内存访问模式优化。Nano-vLLM帮助我们理解了概念模型而生产级的实现则是在此概念上的工程深化。4. 调度器推理引擎的“交通指挥中心”如果说PagedAttention是高效管理“停车场”显存的智慧那么调度器就是确保“车辆”计算任务高效通行的“交通指挥中心”。它的决策直接影响了用户的等待时间延迟和系统的整体处理能力吞吐量。4.1 调度器的核心决策批处理Batching策略在LLM推理中最基本的调度决策就是何时将哪些请求的下一步计算组合成一个批次Batch一起送入GPU执行不同的策略在延迟和吞吐上有着截然不同的表现。静态批处理Static Batching预先确定一个批次大小集齐这么多请求后再开始处理。处理过程中批次成员不变直到所有请求都完成。这是最简单的方式吞吐量高但延迟也最高因为快的请求必须等待慢的请求。动态批处理Dynamic Batching这是vLLM等现代引擎采用的主流策略。调度器持续监控系统中的请求状态。每当GPU完成上一批计算调度器就从所有准备好计算的请求中选择一部分组成下一个批次。一个请求“准备好”意味着它已经收到了用户输入对于首轮或者已经生成了上一个token对于后续轮次。这种方式能更好地平衡延迟和吞吐。在动态批处理中又有一个关键优化连续批处理Continuous Batching 或 Iteration-Level Batching。这是vLLM调度器的精髓。传统动态批处理在请求生成序列长度不同时短序列完成后其GPU资源会闲置等待长序列完成批次才能解散这就是“填充”Padding带来的浪费。连续批处理则允许在每次模型前向传播迭代后就重新组批。4.2 vLLM的连续批处理是如何工作的让我们模拟一个场景看连续批处理如何提升效率时刻 T0请求A输入5个token要求生成10个token和请求B输入3个token要求生成15个token同时到达。调度器将它们放入同一个批次进行预填充Prefill阶段的计算处理所有输入token。假设A需要生成第1个tokenB需要生成第1个token。时刻 T1GPU完成计算A和B都生成了它们的第1个token。调度器更新它们的状态。时刻 T2请求C到达输入4个token要求生成5个token。此时A和B正准备生成下一个token。调度器发现GPU空闲立即将A、B、C组合成一个新批次。注意这个批次里包含了处于解码Decode阶段的A和B它们需要基于已有的KV Cache生成下一个token以及处于预填充阶段的C它需要处理全新的输入token。vLLM的调度器需要智能地处理这种“预填充”和“解码”混合的批次。时刻 T3GPU计算完成。A生成了第2个tokenB生成了第2个tokenC完成了预填充并生成了第1个token。请求C进入解码阶段。时刻 T4请求A完成了它所有的10个token生成从当前批次中退出。调度器释放A占用的KV Cache块如果未被共享。此时批次中只剩下B和C它们继续生成下一个token。同时可能有新的请求D加入...这个过程就像一条流水线请求可以随时加入当有空闲“座位”且其输入就绪时也可以在其任务完成后随时离开而不会阻塞流水线上的其他请求。这极大地提高了GPU的利用率尤其是当请求的生成长度差异很大时。4.3 调度中的关键数据结构与状态管理为了实现这种灵活的调度引擎需要维护每个请求的丰富状态信息。在Nano-vLLM的简化实现中你可能会看到class RequestState: PENDING pending # 等待输入 READY ready # 输入就绪可加入批次 RUNNING running # 正在当前批次中计算 WAITING waiting # 已计算完一轮等待下一轮调度 FINISHED finished # 生成完成 class SchedulingRequest: def __init__(self, request_id, prompt, max_tokens): self.request_id request_id self.prompt prompt self.max_tokens max_tokens self.output_tokens [] self.generated_length 0 self.state RequestState.PENDING # 与PagedAttention相关的状态 self.block_table [] self.last_token_position 0 # 调度优先级可选可用于实现公平排队、优先级队列等 self.priority 0 self.arrival_time time.time()调度器Scheduler的核心循环大致如下class Scheduler: def schedule(self): 调度决策决定下一个批次包含哪些请求 batch_requests [] # 策略1: 优先调度处于READY状态的请求 ready_requests [req for req in self.requests if req.state RequestState.READY] # 策略2: 可能考虑请求的优先级、等待时间、预估计算成本等 sorted_requests self._sort_requests(ready_requests) # 策略3: 进行批次打包考虑GPU内存容量块数量和计算效率 for req in sorted_requests: if self._can_fit_in_batch(batch_requests, req): batch_requests.append(req) req.state RequestState.RUNNING else: break # 当前批次已满 return batch_requests def update_after_computation(self, batch_requests): GPU计算完成后更新请求状态 for req in batch_requests: req.generated_length 1 req.output_tokens.append(new_token) if req.generated_length req.max_tokens: req.state RequestState.FINISHED self._free_blocks_for_request(req) # 释放其KV Cache块 else: req.state RequestState.READY # 准备下一轮生成实操心得在生产环境中调度策略远比上述伪代码复杂。例如需要处理“抢占式调度”高优先级请求打断低优先级请求、“非均匀请求”输入长度差异巨大以及“延迟保障”确保某些请求的最大响应时间。vLLM的调度器实现了多种策略并且其与PagedAttention的深度集成使得基于块数量而非token数量来快速判断“是否能加入批次”成为可能这大大提高了调度决策的效率。5. 实际对比vLLM vs. 其他推理方案理解了vLLM的核心机制后我们将其与社区中其他常见的LLM推理方案进行对比就能更清楚地看到各自的定位与优劣。这也能帮助我们回答“我该选择哪个”的问题。5.1 vLLM vs. 原生PyTorch/Transformers这是最直接的对比。使用transformers库的pipeline或直接调用model.generate()是最简单的方式。优势原生极简部署几行代码即可运行无需额外依赖。灵活性最高可以完全控制模型结构、生成参数方便调试和研究。生态完整无缝接入Hugging Face Hub上的所有模型。劣势原生显存效率低KV Cache管理粗放无法处理长序列或大批次极易OOMOut Of Memory。无高级调度通常只能实现静态批处理吞吐量和延迟难以兼顾。计算未优化缺少针对Transformer解码的特定算子融合优化。结论原生方式适用于原型验证、小模型本地测试、以及需要极致灵活性的研究场景。一旦涉及性能、效率或在线服务就需要更专业的引擎。5.2 vLLM vs. Hugging Face TGI (Text Generation Inference)TGI是Hugging Face官方推出的推理引擎同样以高性能为目标。相似点两者都支持动态批处理、Token流式输出、并针对Transformer解码做了优化。核心差异核心技术vLLM的PagedAttention是其独一无二的招牌在显存管理和共享方面理论优势明显。TGI使用不同的内存管理策略如更传统的缓存管理。模型支持TGI与Hugging Face生态绑定更深对各类Transformer变体模型的支持可能更迅速、更稳定。vLLM也在快速跟进但早期对某些冷门模型架构的支持可能需社区适配。部署与运维TGI提供了开箱即用的Docker镜像和更简单的REST API与Hugging Face的模型库集成极佳部署体验可能更顺畅。vLLM的部署同样简单但更偏向于作为一个高性能库集成到自己的服务中。量化支持两者都支持主流量化方案GPTQ, AWQ等但具体集成方式和易用性有差别。结论如果你需要极致的内存效率和吞吐量尤其是在处理超长上下文或高并发场景时vLLM的PagedAttention可能带来显著优势。如果你追求与Hugging Face生态的无缝集成、快速的模型支持更新和更简单的部署流程TGI是绝佳选择。许多基准测试显示两者在典型场景下性能互有胜负选择常取决于具体模型、硬件和负载特征。5.3 vLLM vs. TensorRT-LLMTensorRT-LLM是NVIDIA推出的闭源推理优化库基于TensorRT。优势TensorRT-LLM极致性能作为硬件厂商的官方工具它能对NVIDIA GPU尤其是最新架构如Hopper进行最深度的优化包括利用新的硬件特性如FP8精度、Transformer Engine通常能获得最高的单卡性能。强大的量化与优化与NVIDIA的量化工具链如AQT结合紧密量化后性能损失小。劣势TensorRT-LLM易用性与灵活性需要将模型编译成特定的引擎格式过程相对复杂调试困难。模型结构支持取决于官方插件对新模型的支持有延迟。调度与多租户其调度能力相比vLLM可能不那么突出更侧重于单批次、单请求的极致性能。在多用户、动态请求的场景下需要上层框架如Triton Inference Server配合。生态锁定紧密绑定NVIDIA硬件和软件栈。结论TensorRT-LLM是追求在特定NVIDIA硬件上达到绝对峰值性能的终极选择适合对延迟和吞吐有极端要求的封闭部署场景。vLLM则在通用性、易用性、动态调度和多租户支持上更胜一筹更适合灵活的云服务或研究环境。5.4 vLLM vs. OllamaOllama是一个专注于本地运行大模型的工具以其极简的安装和使用体验著称。本质区别Ollama更像一个开箱即用的模型运行器和管理器它底层可能使用了已有的推理引擎如llama.cpp并封装了模型下载、版本管理等功能。vLLM则是一个底层的推理引擎库/服务器。适用场景Ollama适合个人开发者、初学者在本地笔记本电脑上快速体验和运行大模型特别是经过量化的模型它的首要目标是易用性和低资源占用。vLLM适合构建生产级API服务、进行严肃的性能测试和优化需要精细控制推理的各个环节。结论两者不在一个赛道。如果你想“五分钟内在Mac上跑通一个模型”选Ollama。如果你想“搭建一个能同时稳定服务上百个用户的长文本生成API”选vLLM。特性/方案vLLM原生 PyTorchTGITensorRT-LLMOllama核心优势PagedAttention, 高显存利用率优秀调度极致灵活易调试HF生态集成部署简单NVIDIA硬件极致性能本地运行极简体验显存效率极高(块化管理共享)低高高 (依赖量化)高 (侧重量化模型)调度能力强(连续批处理)无/弱强中等 (需外部配合)弱性能吞吐很高低高极高(特定硬件)中等 (本地适用)易用性高极高高低极高适用场景生产API高并发长文本研究原型调试生产APIHF模型快速部署超高性能封闭部署个人本地体验与开发选择哪一个最终取决于你的具体需求是追求极致的性能还是极致的易用或是极致的灵活与控制力。对于大多数希望从原型迈向服务的团队而言vLLM和TGI因其在性能、功能和易用性上的平衡成为了最主流的选择。而通过Nano-vLLM理解其内核能让你在使用这些强大工具时更加得心应手也更能诊断和解决可能遇到的问题。