现在位置: 首页 / 大语言模型 / 正文

以 Nano-vLLM 为例，深入理解 LLM 推理引擎（Part 2）

在大语言模型（LLM）的应用浪潮中，推理性能直接决定了用户体验与系统成本。当我们向模型输入一段提示词并等待其逐字生成回复时，背后涉及的是千亿级参数神经网络的复杂计算过程。许多开发者往往将推理引擎视为一个黑盒，仅关注API的调用，却忽视了其内部如何通过重写模型代码与深度优化底层计算逻辑，将静态权重转化为高效的智能输出。理解这一过程，不仅有助于排查性能瓶颈，更是构建高性能AI应用的基础。

本文将以简化版的推理引擎架构为例，深入剖析模型内部的运作机制。我们将重点探讨Token如何转化为高维向量、解码层（Decoder Layer）内部注意力机制与多层感知机（MLP）的具体运算流程、KV Cache在GPU显存中的物理存储策略，以及张量并行（Tensor Parallelism）如何将巨大的计算任务拆解至多个GPU协同处理。通过拆解从Prompt进入系统到文本最终生成的全链路，读者将建立起对LLM推理过程的完整认知，为后续的性能调优与架构设计打下坚实理论基础。

模型的本质构成与推理引擎的代码重构

提到“大语言模型”，大多数人的第一反应是那些体积庞大的权重文件——动辄数十亿甚至上千亿的参数集合。然而，一个真正能够在生产环境中运行并执行推理任务的模型系统，并非仅仅由权重文件组成，而是由三个核心部分协同工作构成的有机整体。

首先是词汇表（Vocabulary），这是一个静态的映射表，负责在人类可读的自然语言文本与模型能够处理的数值表征之间进行双向翻译。它定义了Token（词元）与其唯一ID之间的对应关系，是模型理解语言的基石。其次是权重（Weights），这是模型在大规模训练过程中学习到的参数集合，代表了模型所积累的“知识”。例如，一个7B（70亿参数）模型，意味着它拥有70亿个可调整的浮点数，这些数值决定了模型对语言规律的理解深度。最后，也是最容易被忽视的部分，是运行时代码（Runtime Code）。这部分代码定义了如何使用上述权重，将输入数据一步步转换为输出结果的执行逻辑。这才是真正加载到GPU上并驱动硬件进行矩阵运算的核心程序。

为何推理引擎需要独立实现模型代码

开发者可能会产生疑问：既然模型提供方通常开源了权重文件，为何不直接提供通用的运行时代码？事实上，虽然Hugging Face等平台提供了标准的模型实现，但这些代码往往侧重于通用性和易用性，而非极致的推理性能。标准的Runtime Code通常难以满足生产环境对延迟和吞吐量的严苛要求。

推理场景与训练场景有着本质区别。训练过程注重梯度的精确计算和分布式同步，而推理过程则更关注单次前向传播的速度和显存的高效利用。此外，不同的硬件架构（如NVIDIA A100 vs. 消费级RTX 4090）和精度需求（FP16, BF16, INT4量化）也需要针对性的代码优化。一套在A100集群上表现良好的训练代码，直接移植到单卡推理场景中，可能会因为内核启动开销、内存访问模式不佳等原因导致性能大幅下降。

这正是以vLLM为代表的高性能推理引擎选择重新实现模型代码（Model Code）的根本原因。完整的推理引擎代码库通常包含了对数十种主流模型架构（如Qwen、LLaMA、DeepSeek、Mistral等）的高度优化实现。这些实现针对特定的硬件指令集进行了定制，引入了诸如PagedAttention、连续批处理等高级特性。虽然简化版示例仅支持单一模型架构，但其背后的工程模式——即解耦模型定义与推理调度，并对计算内核进行深度优化——是通用且值得借鉴的最佳实践。

Token在模型中的完整流转管道

为了清晰展示数据在模型内部的流动路径，我们需要追踪一个Token从输入到输出的完整生命周期。这个过程可以大致划分为三个关键阶段：嵌入（Embedding）、解码层处理（Decode Layers）以及语言模型头输出（LM Head）。

Embedding层：从离散Token到连续向量

推理旅程始于嵌入层。当用户输入的文本被分词器（Tokenizer）处理后，会生成一串Token ID序列。对于模型而言，这些ID仅仅是整数索引，例如ID 1547。嵌入层的作用就是在一个巨大的查找表中检索该ID对应的向量表示。

在典型的Qwen模型架构中，每个Token ID会被映射为一个4096维的浮点数向量，这个向量被称为隐藏状态（Hidden State）。它是模型对该Token语义信息的初步内部表征。之所以选择4096维，是在表达能力与计算成本之间经过精密权衡的结果。更高的维度能够捕捉更细微的语义差别和上下文关联，但同时也意味着更大的矩阵运算量和显存占用。嵌入层本质上是一个查找操作，它将离散的符号空间映射到了连续的向量空间，使得后续的数学运算成为可能。

Decode Layers：特征提取与语义深化

随后，初始的隐藏状态会流经一叠解码层。以Nano-vLLM支持的Qwen模型为例，其包含24个相同的解码层堆叠而成。每一层都执行相同的结构操作，但使用各自独立的学习权重。这种层级结构允许模型逐层对表征进行精细化加工。

我们可以这样理解这一过程：底层可能更多关注局部的句法结构和词性关系，中层开始捕捉短语级别的语义含义，而高层则处理更抽象的事实知识和长距离依赖关系。值得注意的是，每一层具体学到了什么特征，并非由工程师预先硬编码，而是在大规模训练过程中自然涌现的结果。

这里的一个关键特性是形状一致性。无论经过多少层处理，每一层接收的输入隐藏状态和输出的隐藏状态，其维度始终保持不变（例如均为4096维）。这种统一性使得深层神经网络可以通过简单的堆叠方式构建，极大地简化了模型架构的设计与实现。

LM Head：从向量空间回归概率分布

经过所有解码层的深度处理后，最终的隐藏状态包含了丰富的上下文信息。此时，模型需要将这些高维向量转换回人类可读的文本。这一任务由语言模型头（LM Head）完成。

LM Head的功能可以视为嵌入过程的逆向操作。它通过一个线性变换层，将4096维的隐藏状态投影回词汇表的大小（例如15万维）。输出的结果称为Logits，即词汇表中每个可能下一个Token的非标准化打分。这些Logits随后会被送入采样算法（如Softmax、Top-K、Top-P采样），转化为概率分布，并最终选出概率最高的Token作为模型的输出。至此，一个完整的从“理解”到“生成”的闭环得以完成。

解码层内部构造详解：注意力机制与MLP

解码层是大语言模型的核心计算单元，其内部主要包含两大模块：多头注意力机制（Multi-Head Attention, MHA）和多层感知机（Multi-Layer Perceptron, MLP）。这两个模块交替工作，共同完成了对上下文信息的整合与特征提炼。

多头注意力机制：上下文感知的核心

注意力机制赋予了模型“关注”序列中其他Token的能力，使其能够理解上下文依赖关系。然而，现代LLM并不使用简单的单头注意力，而是采用多头注意力机制，将计算任务拆分到多个并行的“头（Head）”中，以捕捉不同子空间的语义特征。

在Qwen模型中，共有32个注意力头。每个头处理128维的数据切片（$32 \times 128 = 4096$，恰好等于隐藏状态的总维度）。这不仅仅是简单的维度切分，每个头都会通过独立的线性投影矩阵，将完整的4096维输入变换为该头特有的128维表征。这种设计允许不同的头专注于不同类型的关系，例如有的头可能关注语法结构，有的关注指代关系，有的关注实体关联。

可以将这一过程想象成一个拥有32个工作站的工厂流水线。每个工作站接收相同的原材料（完整的隐藏状态），但使用不同的工具（投影权重）将其塑造成特定的形状。虽然我们无法确切知道每个头具体学到了什么，但这种并行机制显著增强了模型的表达能力。

在每个头内部，模型会计算当前Token与序列中所有先前Token的相关性分数（Attention Scores）。这正是模型理解诸如“The cat sat on the mat. It was comfortable.”中“It”指代“the cat”的关键所在。所有头的输出随后被拼接（Concatenate）并通过一个输出投影矩阵，重新映射回4096维，形成该层的注意力输出。

MLP：非线性特征变换与信息增强

紧随注意力机制之后的是多层感知机（MLP）模块。与注意力机制不同，MLP不具备跨Token的交互能力，它独立地处理每个Token的隐藏状态。其核心作用是对注意力机制提取的特征进行进一步的非线性变换和信息增强。

MLP的典型结构包含两个线性层和一个激活函数。首先，它将4096维的隐藏状态扩展到一个更大的中间维度（在Qwen模型中为11008维），这一步被称为“上采样（Upscale）”。接着，应用非线性激活函数（如SiLU或GeLU），引入非线性因素，使模型能够拟合复杂的函数关系。最后，再通过第二个线性层将维度压缩回4096维。

为什么要进行这样的“扩展-压缩”操作？可以将其类比为图像处理中的分辨率提升。4096维的隐藏状态可能包含了一些压缩的信息，扩展到11008维创造了更多的“空间”，让模型能够添加细节和微调特征。这些细节由MLP的学习权重决定，反映了模型在训练中学到的特定知识模式。随后的压缩过程则是对这些经过丰富后的表征进行提炼，去除冗余信息，保留最核心的语义特征，传递给下一层解码层。这种瓶状结构（Bottle-neck Structure）在保证信息容量的同时，有效地控制了计算复杂度。

4.3 Dense 架构与 MoE 架构的深度博弈

在前文对 MLP 模块的解析中，我们默认采用的是 Dense（稠密）架构，即每一个输入 Token 都会激活并流经模型中所有的神经元参数。然而，随着模型规模的指数级增长，另一种名为 Mixture of Experts (MoE，混合专家) 的架构逐渐成为了主流选择。MoE 的核心思想并非使用一个巨大的全连接层，而是将其拆分为多个小型的、独立的“专家”网络（Expert Networks），例如将一个大 MLP 替换为 8 个较小的 Expert MLPs。

image.png

在 MoE 机制下，引入了一个关键的组件——路由网络（Router Network），也常被称为门控网络（Gating Network）。对于每一个输入的 Hidden State，路由网络会动态计算其属于各个专家的权重，并据此决定由哪几个专家来处理该 Token。以常见的 Top-2 路由策略为例，即使模型拥有 8 个专家，对于任意给定的 Token，通常只激活其中得分最高的 2 个专家进行前向传播，其余专家则保持静默。这种稀疏激活机制是 MoE 高效性的根源。

许多初学者容易陷入一种拟人化的误区，认为不同的专家会被训练成专门处理数学、代码或特定语言领域的“专科医生”。事实上，这种专业分工并非人为预设，而是从海量数据训练中涌现（Emergent）出来的特性。虽然我们在事后分析中可能发现某个专家对代码逻辑更敏感，但在大多数情况下，专家之间的界限是模糊且高度重叠的。我们很难清晰地界定 Expert 3 与 Expert 5 在具体功能上的绝对差异，它们更多是共同协作以拟合复杂的数据分布。

引入 MoE 的核心动机在于提升计算效率（Computational Efficiency），而非直接追求单次推理的输出质量上限。 通过稀疏激活，MoE 允许模型拥有极其庞大的总参数量（所有专家参数之和），但在推理时，每个 Token 仅消耗极少部分的计算资源。这意味着，在相同的算力预算下，MoE 模型可以比 Dense 模型大出一个数量级，从而具备更强的知识容纳能力和泛化潜力，同时保持较低的推理延迟。

当然，这种设计并非没有代价，这是一种典型的工程权衡。在总参数量相同的前提下，Dense 模型往往能产生比 MoE 模型更稳定、高质量的输出，因为 Dense 模型对每个 Token 都利用了全部信息进行精细化处理。然而，当模型规模扩展到数千亿参数时，Dense 模型的训练成本和推理显存需求变得令人望而却步。MoE 的价值在于它打破了算力的硬约束，以牺牲少量的单参数效率为代价，换取了模型规模的可扩展性（Scalability）和实际可训练性（Practical Trainability）。

05 KV Cache：数据平面的物理实现

在第一部分中，我们将 Block Manager 视为 KV Cache 的控制平面，主要负责在 CPU 内存中追踪逻辑块的分配状态、引用计数以及碎片整理。现在，我们将视角下沉到数据平面，深入探讨 KV Cache 究竟是如何物理存储在 GPU 显存中的，以及底层内核如何高效地访问这些数据。理解这一层对于优化显存利用率和降低内存带宽瓶颈至关重要。

5.1 缓存内容的本质：K 与 V 向量

在 Transformer 的 Attention 计算机制中，每个 Token 经过线性变换后会产生三个向量：Query (Q)、Key (K) 和 Value (V)。在自回归生成过程中，当前 Token 的 Q 需要与之前所有 Token 的 K 进行点积运算以计算注意力分数，并与 V 进行加权求和。由于之前 Token 的 K 和 V 在后续步骤中不会改变，为了避免在每一步 Decode 时重复计算历史 Token 的 K/V 值，我们将它们持久化存储在显存中，这就是 KV Cache 的由来。

5.2 GPU 显存中的多维物理布局

在 GPU 显存中，KV Cache 并非简单的线性数组，而是一个精心组织的多维张量结构，旨在最大化内存访问的局部性和并行度。其物理布局通常包含以下四个关键维度：

  • Block 维度：这与 Block Manager 中的逻辑块概念相对应。为了减少内存碎片并提高分配效率，显存通常以固定大小的块（例如每块容纳 256 个 Token）为单位进行预分配。
  • Layer 维度：Transformer 模型由多层 Decoder 组成（如 24 层或 80 层）。由于每一层的注意力机制是独立计算的，因此每一层都需要维护自己独立的 KV 缓存空间。
  • K/V 维度：在每一层内部，Key 和 Value 是两个独立的张量。因此，物理上需要为每层分别分配 K 缓存区和 V 缓存区。
  • Token/Head 维度：在每个块内部，需要为每个 Token 的每个注意力头（Attention Head）预留具体的向量空间。

因此，Block Manager 中的一个逻辑块索引，在 GPU 物理显存中实际上映射到了 $L \times 2$ 个连续的内存区域（其中 $L$ 为层数）。例如，在一个 24 层的模型中，一个逻辑块对应着 48 个物理缓存片段（24 层 $\times$ K/V）。这种映射关系使得引擎能够通过简单的算术运算快速定位到任意层、任意 Token 的 K 或 V 数据。

5.3 基于 Triton 的高效内核访问

传统的 CUDA 编程虽然性能强大，但编写和维护复杂的内存访问内核极为繁琐。现代推理引擎（如 vLLM 及其简化版实现）倾向于使用 Triton 这种高级 GPU 编程语言。Triton 允许开发者以类似 Python 的方式编写内核，编译器会自动将其优化为高效的 CUDA PTX 代码。

@triton.jit
def load_kv_cache(
    kv_cache_ptr,   # 指向显存中 KV Cache 基地址的指针
    block_tables,   # 逻辑块到物理块索引的映射表
    seq_len,        # 当前序列长度
    stride_l,       # Layer 维度的步长
    stride_h,       # Head 维度的步长
    BLOCK_SIZE: tl.constexpr
):
    # 计算当前处理的 token 索引和 layer 索引
    pid = tl.program_id(axis=0)
    layer_idx = pid // num_heads
    head_idx = pid % num_heads

    # 根据逻辑块表查找物理块位置
    block_idx = tl.load(block_tables + token_idx // BLOCK_SIZE)
    offset_in_block = token_idx % BLOCK_SIZE

    # 计算最终显存地址并加载数据
    ptr = kv_cache_ptr + block_idx * stride_b + layer_idx * stride_l + head_idx * stride_h + offset_in_block
    kv_data = tl.load(ptr)
    return kv_data

上述代码展示了 Triton 内核如何封装底层的地址计算逻辑。block_tables 提供了从逻辑序列到物理显存块的映射，内核通过 tl.load 指令直接从 GPU 全局内存中读取数据。这种方式不仅代码可读性高，而且 Triton 编译器能够自动处理内存合并访问（Coalesced Access）和寄存器分配，从而在保证灵活性的同时达到接近手写 CUDA 的性能。

06 张量并行（Tensor Parallelism）：计算层面的拆解

在第一部分中，我们介绍了张量并行（TP）的通信拓扑，以及主进程如何通过共享内存广播执行命令。本节将聚焦于计算层面，详细解析 模型权重和计算任务是如何被拆分到多个 GPU 上协同工作的。张量并行的核心在于将大型矩阵运算切分为小块，分布在不同的设备上并行执行，最后通过通信操作合并结果。

6.1 Attention 机制中的头部并行

以 TP=2（两张 GPU 并行）为例，当隐藏状态向量进入 Attention 模块时，并行策略如下：

image.png

  1. 输入复制：两张 GPU 均接收完整的隐藏状态向量（例如 4096 维）。在此阶段，输入数据并未被拆分，确保每张卡都有完整的上下文信息用于计算 Query。
  2. 头部切片（Head Splitting）：Attention 的计算压力主要来自于多头注意力机制。系统将注意力头（Attention Heads）均匀分配给各张 GPU。例如，若共有 32 个头，GPU 0 负责计算 Head 0-15，GPU 1 负责计算 Head 16-31。每张卡只持有对应头的 Q、K、V 投影权重。
  3. 局部计算：每张 GPU 独立计算其负责的注意力头的输出。此时，GPU 0 的输出仅包含前半部分头的信息，GPU 1 仅包含后半部分。
  4. All-Reduce 合并：为了得到完整的 Attention 输出向量，两张 GPU 执行 All-Reduce 操作。它们交换各自的部分结果并进行求和。最终，两张 GPU 上都得到了完整的、融合了所有头信息的 Attention 输出向量，供下一层 MLP 使用。

这种并行方式的关键在于：并行发生在 Head 维度，而非 Sequence 或 Hidden Dimension 维度。这保证了每个 Token 的全局上下文一致性，同时大幅降低了单卡的计算负载。

6.2 MLP 层中的列并行与行并行

MLP 层的并行策略略有不同，通常采用经典的 Megatron-LM 模式，结合列并行（Column Parallelism）和行并行（Row Parallelism）：

image.png

  1. 输入复制：同样，每张 GPU 接收完整的上一层输出向量。
  2. 列并行（Up/Gate Projection）：MLP 的第一层通常是将隐藏维度投影到更大的中间维度（例如从 4096 扩展到 11008）。这个巨大的权重矩阵按列切分。若 TP=2，每张 GPU 只存储并计算一半的中间维度（5504 维）。此时不需要通信，因为各卡独立计算各自的部分。
  3. 非线性激活：每张 GPU 对其计算出的部分中间向量应用激活函数（如 SiLU 或 GeLU）。
  4. 行并行（Down Projection）与 All-Reduce：MLP 的第二层将中间维度投影回隐藏维度。这个权重矩阵按行切分。每张 GPU 计算部分输出后，必须通过 All-Reduce 将结果相加，才能还原出完整的隐藏状态向量。这一步确保了数据流的完整性，以便传递给下一个 Transformer 层。

6.3 通信开销与适用场景

张量并行并非银弹，其最大的瓶颈在于 通信开销（Communication Overhead）。All-Reduce 操作需要在 GPU 之间同步大量数据，这会引入显著的延迟。因此，TP 的有效性高度依赖于互联带宽。

  • 高速互联场景：在单机多卡环境中，利用 NVLink 或 NVSwitch 提供的高带宽（数百 GB/s），TP 带来的计算加速远大于通信延迟，效果显著。
  • 跨机互联场景：若通过以太网或 InfiniBand 跨节点并行，通信延迟将成为主导因素，导致加速比急剧下降甚至负优化。

此外，TP 的主要优势在于显存节省。通过 TP=8，每张 GPU 只需存储模型权重的 1/8。这使得在单卡显存有限的情况下，运行参数量远超单卡容量的超大模型成为可能。它是解决“显存墙”问题的关键技术手段。

07 思考：架构设计中的权衡与哲学

在深入剖析了底层机制后，我们有必要跳出代码细节，从系统设计的角度审视一些常见的架构抉择。理解这些权衡有助于我们在面对具体业务场景时做出更明智的技术选型。

7.1 深度（Layers）与宽度（Heads）的平衡

网络层数（Depth）主要决定了模型的推理深度和抽象能力。 每一层 Transformer 都可以看作是对隐藏状态的一次迭代细化，层数越多，模型越能捕捉长距离依赖和复杂的逻辑链条。而 注意力头数（Width/Heads）则提供了并行处理不同语义关系的“视角”，更多的头意味着模型能同时关注语法、语义、指代等多种特征。

是否存在极端的优化空间？例如，构建一个“窄而深”（少头多层）的模型以增强逻辑推理，或“宽而浅”（多头少层）的模型以覆盖更广的知识面？研究表明，这种极端偏离均衡比例的架构往往表现不佳。就像人类认知需要广度与深度的结合一样，Transformer 模型也需要在两者之间保持大致均衡的比例（Scaling Laws 暗示了这一点）。

真正撬动模型性能杠杆的，往往不是微调架构比例，而是 训练数据的质量与多样性 以及 训练策略的有效性。在算力固定的前提下，增加高质量数据或使用更先进的优化器，通常比调整层数与头数的比例带来更稳定的收益。

7.2 MoE 流行的深层逻辑：规模即正义

再次回到 MoE 架构，其流行并非源于它在单位参数下的效率优势。相反，一个 70B 参数的 Dense 模型，在同等参数量下通常比 MoE 模型表现更好、更稳定。那么，为什么业界纷纷转向 MoE？

答案在于“可触及的规模上限”。 训练一个 600B 参数的 Dense 模型，所需的算力集群规模和能源消耗是目前大多数机构无法承受的。然而，一个总参数 600B、但每次推理仅激活 50B 参数的 MoE 模型，其训练和推理成本却处于可接受范围内。

这是一种务实的工程哲学：以轻微的单位参数效率损失，换取模型总体规模的数量级跃升。 凭借庞大的总参数量，MoE 模型能够容纳远超 Dense 模型的知识量，从而在综合评测中展现出更强的能力。在当前的技术边界下，规模（Scale）仍然是提升智能水平最确定的路径，而 MoE 则是通往更大规模的必经桥梁。

08 结语

至此，我们完成了一次从微观算子到宏观架构的深度巡礼。回顾整个推理链路：

  • Tokenization 将自然语言离散化为机器可理解的 ID 序列；
  • Embedding 将这些 ID 映射为高维语义向量；
  • Decoder Layers 通过 Attention 机制捕捉上下文关联，通过 MLP 整合内部知识，层层细化隐藏状态；
  • KV Cache 作为记忆中枢，避免了历史计算的冗余，极大提升了生成速度；
  • LM Head 将最终的语义向量还原为词汇表的概率分布；
  • Sampling 策略从中选取下一个 Token，完成一次迭代；
  • Tensor Parallelism 则在底层支撑起这一切，让超大模型得以在多卡间高效流转。

推理引擎不仅是代码的执行者，更是资源的调度者。它统筹着请求队列、内存分配、并行通信与计算内核，将复杂的数学变换转化为流畅的文字输出。大语言模型的本质，是一个经过海量数据训练的、极其精密的向量变换函数。所谓的“智能”，并非魔法，而是参数规模、数据质量与系统工程巧思共同作用的涌现结果。

无论您是致力于生产环境的性能调优，还是出于好奇探索 AI 的黑盒内部，希望这篇关于推理引擎内部机制的解析，能为您建立起清晰、扎实的技术认知框架。

END

本期互动内容 🍻

❓ 文章提到推理引擎需要重写 Model Code 以适配深度优化（如 Triton Kernel 或 TP 拆分）。在您自己的实践中，有没有遇到过“理论算法可行，但工程落地跑不动”的部署场景？最后是通过裁剪模型、量化还是其他手段妥协突破的？欢迎在评论区分享您的踩坑经验！

(本文经原作者授权，由技术博客编辑团队编译润色。如需转载，请联系获取授权。)

原文链接： neutree.ai/blog/nano-vllm-part-2