Nano-vLLM深度解读（下）

爱听歌的周童鞋

592人浏览 · 2026-04-07 21:06:11

爱听歌的周童鞋 · 2026-04-07 21:06:11 发布

前言

本篇文章继续我们的 nano-vllm 解读，记录下个人学习笔记，和大家一起分享交流😄

Note：一些基础的内容默认大家都会，我们的重点是关注 vLLM 推理框架的核心机制，而对于模型本身的知识例如 nano-vllm 使用的 Qwen 系列模型构建部分我们会跳过，不会讲得那么详细。

reference：https://github.com/GeeeekExplorer/nano-vllm

reference：https://www.bilibili.com/video/BV18EfdB9E2Y

reference：https://chatgpt.com/

1. 本期内容概览

在本篇文章中，我们将继续对 nano-vLLM 项目进行深入解读。nano-vLLM 是当前非常流行的开源推理引擎 vLLM 的一个迷你实现版本。在上一篇文章 Nano-vLLM 深度解读（上）中，我们已经介绍了该项目的整体定位，以及项目作者的背景信息。

整个系列文章分为 上下两篇，核心目标是：通过完整梳理 nano-vLLM 的源码结构，总结出其整体架构，并逐步拆解其中涉及的关键模块与核心概念。最终希望和大家一起来理解：

推理引擎究竟在做什么
它需要解决哪些核心问题
在实现过程中使用了哪些关键技术

通过对这些问题的逐步分析，我们最终整理出了如下所示的整体架构图：

在这里插入图片描述

OK，那在上一篇文章中，我们重点介绍了推理引擎中的 调度策略（Scheduler） 以及相关功能特性的实现方式，并对其进行了逐步拆解与讲解。大家感兴趣的可以先阅读上一篇文章，以获得更完整的背景理解。

在上一篇文章的结尾，我们刻意留下了一个 “黑盒模块” — 即下图中的灰色部分 Model & KV Cache：

我们提到，所有发生在 GPU 上的核心计算过程，以及模型推理的关键逻辑，实际上都集中在这个黑盒模块之中，这个模块正是我们所说的 模型与 KV 缓存（Model & KV Cache），也是上一篇文章最后留下的悬念。

而在本文中，我们将正式打开这个“黑盒”，深入分析其中的实现细节，包括：

模型内部的计算流程
KV Cache 的实现机制
Tensor Parallelism（TP，张量并行）等核心功能在底层数据结构和计算层面的实现方式

换句话说，本篇文章的重点是：从系统架构层面进一步下沉到模型执行与 GPU 计算的核心实现部分。

在这里插入图片描述

如上图所示，当我们将这一黑盒模块进一步展开之后，就得到了本期文章将要涉及的完整架构结构。可以看到，原本在上篇文章中作为整体模块存在的 Model & KV Cache，在这里被进一步拆解成多个具体组件，而这些组件也正是本篇文章将要逐一分析的内容。

以上就是我们的一个大体的讲解思路。

OK，在正式进入架构图的详细讲解之前，我们先来看一下 本期内容的整体大纲：

回顾本期定位
- 上期内容回顾
- 本期目标：展开 Model & KV Cache 黑盒
模型基础
- 什么是模型，为什么推理引擎要「实现模型」
- 模型整体流水线：Embedding → Decoder Layers × N → LM Head
- 现代 LLM 的优化手段 （本次内容不展开）
  - RoPE：更好的位置编码，支持长文本
  - RMSNorm：更高效的归一化方式
  - GQA：减少 KV 头数，降低显存占用
  - Flash Attention：分块计算 Attention，提升速度并降低显存
深入 Layer 内部
- Multi-Head Attention
  - 类似 “圆桌会议”，每个 token 可以获取其他 token 的信息
- MLP（传统 Dense 模型）
  - 对当前 token 表示进行进一步的非线性变换
- MoE MLP
  - 通过专家路由机制优化计算量
与 KV cache 的数据层结合
- KV Cache 仅在 Attention 阶段 参与计算
- Block Manager 与 KV Cache 的数据映射关系
- 使用 Triton Kernel 进行计算优化
实现 Tensor Parallel
- Attention 与 MLP 中的并行计算
- 理解 Column Parallel / Row Parallel 及其通信需求
思考问题
- 是否可以通过设计 Layer 数量与 Head 数量 来定义模型风格？
- 为什么 MoE 模型 正在变得越来越流行？

接下来，我们首先会从 模型基础 开始讲起，在这一部分中，我们会重新讨论一个看似简单、但实际上非常关键的问题：什么是模型？

很多人第一次接触推理引擎时都会有一个疑问：既然我们已经加载了一个 巨大的模型权重文件，为什么推理引擎内部仍然需要实现这么多模型相关的代码？为什么不能直接使用权重进行推理呢？🤔

理解这个问题，其实就是在理解 模型本身到底是什么，这一点，也是我们接下来要重点说明的内容。

在理解模型之后，我们会进一步来看 现代 Transformer 模型的整体计算流水线，从结构上来看，一个典型的大语言模型通常可以抽象为 Embedding → Decoder Layers × N → LM Head 三部分，其中：

Embedding：将输入 token 转换为向量表示
Decoder Layers × N：模型最核心的计算部分
LM Head：将隐藏状态映射回词表空间，得到 logits

在整个流水线中，真正承担主要计算任务的其实是中间的 Decoder Layers，而前后的 Embedding 与 LM Head 本质上只是比较标准的向量变换过程，这一点我们一会就会看到。

接下来，我们还需要提前说明下 本期内容不会展开的部分 即 现代 LLM 的一些模型优化技术。这些技术实际上在 nano-vllm 的代码中是存在的，因为 nano-vllm 使用的演示模型是 Qwen 系列模型，因此在模型实现中也包含了 Qwen 所采用的一些主流优化方法，例如：

RoPE：改进的位置编码方式，用于支持更长上下文
RMSNorm：比 LayerNorm 更轻量的归一化方式
GQA（Grouped Query Attention）：减少 KV 头数量，从而节省显存
FlashAttention：高效的 Attention 计算算法

其中有些优化是 在 nano-vLLM 仓库中直接实现的，而像 FlashAttention 则是通过 调用第三方库来完成的。不过，在本篇文章中我们不会深入分析这些技术，主要原因是：这些优化更多是 具体模型实现层面的细节，而不是 推理引擎架构本身的核心问题。

换句话说，这些技术虽然在现代 LLM 中非常常见，但它们通常是 与具体模型实现强绑定的，例如：

Qwen 可能使用 RoPE + GQA
DeepSeek 可能采用不同的 Attention 变体
Mistral 可能又有不同的设计选择

虽然这些方法本质上都属于 现代 LLM 的优化手段，但它们的具体实现往往会随着模型架构的不同而有所变化。因此，从 推理引擎的角度 来看，这部分内容并不是本次分析的重点。

我们在这里先将 nano-vllm 中与这些优化相关的一部分代码 暂时“摘出来”，先不去关注它。这样做的目的，是让我们能够更加专注于 推理引擎如何组织模型计算，以及如何管理 KV Cache 与并行计算 这些更核心的问题。

在讲解完模型基础知识之后，我们将把关注点放在 Decoder Layer 的内部结构 上。

在现代 Transformer 模型中，Decoder Layer 是整个模型计算的核心单元，几乎所有基于 Transformer 架构的语言模型都会采用类似的结构，因此，我们会深入到 Layer 内部，重点理解其中几个关键组成部分。

首先是 Attention 机制，大家可能都听说过 Attention，它是 Transformer 架构中最核心的概念之一，不过，在工程实现层面，推理引擎通常不会只实现一个简单的 Attention，而是会使用 Multi-Head Attention（多头注意力）。

Multi-Head Attention 的核心思想是：将注意力计算拆分成多个 attention head 并行计算，每个 head 可以关注输入序列中的不同信息模式。这样做不仅可以提升模型表达能力，也方便在 GPU 上进行并行计算，最终得到更高质量的输出表示。

在 Attention 之后，Decoder Layer 中还有另一个同样重要的模块：MLP（Feed-Forward Network）。相比 Attention，MLP 在讨论中往往出现得较少，因为它的工程实现相对简单，但事实上，MLP 才是模型中占据大量参数的部分，在许多 Transformer 模型中，大部分权重参数都集中在这一层。

在本次内容中，我们会将 MLP 分成两类来介绍，因为这与模型架构的设计密切相关：

第一类是 传统的 Dense 模型。“Dense” 通常翻译为“ 稠密”，意思是 每个 token 都会经过完整的 MLP 计算路径。这种结构实现简单、计算稳定，也是早期 Transformer 模型中最常见的设计方式。

第二类则是对 Dense MLP 的一种优化形式，也就是近年来非常流行的 MoE（Mixture of Experts）架构。在 MoE 结构中，一个 MLP 层不再只有一条计算路径，而是由 多个 Expert（专家网络）共同组成，通过一个路由机制，每个 token 只会被分配到部分 expert 上进行计算，从而在保持模型容量的同时减少计算量。

在这一部分的分析中，我们也会更加清晰地看到 MoE 实际上作用在 MLP 层，而不是 Attention 层。

很多人在第一次接触 MoE 时，可能会将 “专家” 理解为类似人类专家的概念，例如擅长数学的专家、擅长语言的专家、擅长编程的专家，但从模型角度来看，expert 本质上是一种计算结构，而不是明确的人类知识划分。

这些 expert 是在训练过程中通过优化目标 自然形成的计算分工，它们可能在统计意义上呈现出某种“专业化”的倾向，但这种划分通常很难与人类知识领域进行一一对应。因此，expert 更准确的理解方式是：模型内部通过训练形成的一种概率意义上的计算分工结构。

目前，如何更精确地控制 expert 所学习的知识内容，仍然是一个非常前沿的研究方向。从工程角度来看，MoE 的主要目标其实是 在不显著增加计算成本的前提下，提高模型容量，这也是为什么近年来 MoE 架构开始变得越来越流行的原因之一，后面我们也会进一步讨论这一点。

在理解了 Decoder Layer 的结构之后，我们还需要进一步分析 KV Cache 在数据层面上的实现方式。在上一篇文章中，我们已经介绍过 Block Manager，Block Manager 本质上属于 控制面（control plane） 的组件，它运行在 CPU 上，负责管理序列与 KV Cache block 之间的逻辑映射关系。而真正存储在 GPU 显存中的，则是 KV Cache 本身，这些数据构成了推理过程中的 数据面（data plane）。

因此，一个完整的推理系统实际上需要解决这样一个问题：如何让 CPU 侧的 Block Manager 与 GPU 侧的 KV Cache 在数据结构上建立一致的映射关系，在后续内容中，我们会进一步分析这一部分的实现逻辑。

最后，我们还会介绍推理引擎中的一个重要能力：Tensor Parallel（张量并行）。

Tensor Parallel 可以理解为推理引擎提供的一种 并行执行能力，需要注意的是 并行计算并不是模型天然具备的能力，而是由推理框架实现的。因此，我们也会具体分析推理引擎在实现 Tensor Parallel 时，是如何在不同计算阶段进行拆分的。例如：

在 Attention 阶段 如何进行并行
在 MLP 阶段 如何进行并行

在这一过程中，还会涉及两个非常重要的概念：Column Parallel 和 Row Parallel，这两个概念在第一次接触时往往比较容易混淆，在后面的内容中，我们会结合下面这张图来帮助大家建立更加直观的理解：

通过这张图，我们可以更加清楚地看到 Column Parallel 与 Row Parallel 的计算拆分方式，以及它们所需要的通信模式。

当本篇文章的所有内容讲完之后，我们希望能够引导大家一起思考两个问题：

第一个问题是关于 模型结构设计。在前面的内容中我们提到了两个重要概念：

Layer：Decoder Layer 的层数
Head：Multi-Head Attention 中的 head 数量

那么，这两个参数在模型中的作用分别是什么呢？它们是否能够在一定程度上 定义模型的能力或风格？

第二个问题则与 MoE 架构有关。

为什么 MoE 会被提出？它变得越来越流行的原因是什么？MoE 的目标是提升模型输出质量？还是主要为了 优化计算量？

在后面的分析中，我们也会逐步回答这些问题。

OK，以上就是本篇文章的整体内容大纲。

2. 模型基础

接下来，我们正式进入整个架构图中的第一个部分：模型基础，在这一节中，我们将回答两个问题：什么是模型？一个典型语言模型的计算流水线是什么样的？

架构图第一个部分的结构如下图所示：

首先，我们可以给出一个非常简化但又非常有帮助的表达式：Model = vocab + weights + runtime code。

当我们在日常讨论 “大模型” 时，很多时候关注的其实主要是 权重（weights），例如新闻中经常提到 1B 参数模型、7B 参数模型、70B 参数模型甚至 600B 或 1000B 级别的模型，这里的 B（Billion） 表示模型的参数规模，也就是 权重数量。因此很多人会产生一个非常直观的理解：参数越多，模型能力越强。

这种理解在很多情况下确实是成立的，但如果我们真正从 推理系统的角度 来看，一个模型并不仅仅只有权重。当模型真正运行起来进行推理时，它至少需要以下三个核心要素：

1. vocab（词汇表）

第一部分是 vocabulary（词汇表）。

当模型进行运行阶段时，所有自然语言文本都会被转换为 数字化表示，例如：

"hello world"
      ↓
[15339, 1917]

在模型内部，所有计算都只是在处理这些 token id 对应的数字向量，但在最终输出时，我们仍然需要将这些数字 重新映射回自然语言文本，而这个映射关系，就是由 vocabulary（词汇表） 定义的。

因此，vocab 的作用可以理解为：

定义 token 与自然语言之间的映射关系
决定模型可以表达的基本符号集合

在工程实现上，vocab 通常是一个 完全静态的文件，不会在运行过程中发生变化。

2. weights（模型权重）

第二部分就是我们最熟悉的 模型权重（weights）。

权重是模型在训练过程中学习得到的参数，它决定了模型如何从输入 token 逐步计算出最终的输出结果，这些权重通常会以类似如下形式存在：

pytorch_model.bin
model.safetensors

对于一个 7B 模型来说，这些权重文件通常会达到 十几 GB 的规模，在推理阶段，这些权重也是 完全静态的，不会被修改，因此从工程角度来看 vocab 与 weights 本质上都是静态资源文件。

3. runtime code（运行时逻辑）

第三部分则是很多人在第一次接触推理系统时容易忽略的部分：runtime code（运行时代码）。

即使我们已经拥有了 vocabulary 和 model weights，模型本身仍然 无法直接运行，因为我们还需要一段程序逻辑来完成以下事情：

将 token id 转换为 embedding
按顺序执行每一层 decoder layer
在 Attention 与 MLP 中完成计算
最终通过 LM Head 得到 logits

换句话说，权重只是数据，而 runtime code 才是执行这些计算的程序。在上图中，红色框标出的部分就是 runtime code。

Note：需要说明的是，除了上述核心要素之外，模型推理还可能涉及一些周边功能（如对话模板、工具调用解析器等等）。它们服务于特定场景，不属于推理核心依赖，因此我们在这里就不展开了。

这里就会产生一个很自然的问题：既然模型在推理时需要 runtime code，那么模型发布时直接把 runtime code 一起开源出来，不就可以了吗？🤔

事实上，很多模型在早期发布时确实会提供一份简单的 runtime 实现，例如一个 Python 文件或一组脚本，用来演示如何运行该模型。但随着工程需求的不断增加，人们逐渐发现：由模型发布者维护一份统一的 runtime code，其实是非常困难的。

原因在于，runtime code 的目标是 在特定硬件环境下高效运行模型，而不同场景之间存在大量差异，例如：

推理发生在 训练过程中，还是 生产环境
推理运行在 CPU 还是 GPU
GPU 的具体架构（A100 / H100 / consumer GPU 等）
是否需要 高并发服务
是否需要 分布式推理

在不同场景下，runtime 的设计都会产生不同的工程取舍。

举一个最简单的例子：如果是在 训练过程中 进行推理，那么模型权重通常是 可更新的，runtime code 需要支持梯度计算和参数更新。而在 纯推理场景 下，权重已经完全固定，这时我们就可以移除训练相关逻辑，使用更激进的优化策略最大化推理性能，因此，很难用 一套 runtime code 同时适配所有场景。

正因为如此，在实际工程中逐渐形成了一种分工模式：

模型提供方负责发布：

vocabulary
model weights
一些基础配置文件

而 runtime code 则由不同的推理框架来实现，例如：

HuggingFace Transformers
vLLM
SGLang
llama.cpp
TensorRT-LLM

模型厂商如果希望自己的模型能够被广泛使用，通常需要在这些主流推理框架中 适配对应的模型实现。这也是为什么在 nano-vllm 的仓库中，我们可以看到大量与 Qwen 模型相关的代码，这些代码实际上就是 Qwen 模型在 nano-vLLM 推理框架中的 runtime 实现。如果你查看完整的 vLLM 仓库，会发现其中包含大量不同模型的 runtime 代码，例如 Qwen、DeepSeek、Mistral、LLaMA 等等。

OK，在理解了 Model = vocab + weights + runtime code 之后，我们今天要重点关注的其实就是 在 vLLM（以及 nano-vllm）这样的推理框架中，runtime code 到底包含哪些内容，它是如何实现的？

这也正是前面我们提出的一个关键问题：为什么推理引擎需要 “实现模型”？当这个问题想清楚之后，我们就可以从一个更加宏观的角度，来看一看 模型推理的整体计算流水线：

完整的模型代码如下：

# nanovllm/models/qwen3.py
import torch
from torch import nn
import torch.distributed as dist
from transformers import Qwen3Config

from nanovllm.layers.activation import SiluAndMul
from nanovllm.layers.attention import Attention
from nanovllm.layers.layernorm import RMSNorm
from nanovllm.layers.linear import QKVParallelLinear, MergedColumnParallelLinear, RowParallelLinear
from nanovllm.layers.rotary_embedding import get_rope
from nanovllm.layers.embed_head import VocabParallelEmbedding, ParallelLMHead


class Qwen3Attention(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        num_heads: int,
        num_kv_heads: int,
        max_position: int = 4096 * 32,
        head_dim: int | None = None,
        rms_norm_eps: float = 1e-06,
        qkv_bias: bool = False,
        rope_theta: float = 10000,
        rope_scaling: tuple | None = None,
    ) -> None:
        super().__init__()
        tp_size = dist.get_world_size()
        self.total_num_heads = num_heads
        assert self.total_num_heads % tp_size == 0
        self.num_heads = self.total_num_heads // tp_size
        self.total_num_kv_heads = num_kv_heads
        assert self.total_num_kv_heads % tp_size == 0
        self.num_kv_heads = self.total_num_kv_heads // tp_size
        self.head_dim = head_dim or hidden_size // self.total_num_heads
        self.q_size = self.num_heads * self.head_dim
        self.kv_size = self.num_kv_heads * self.head_dim
        self.scaling = self.head_dim ** -0.5
        self.qkv_bias = qkv_bias

        self.qkv_proj = QKVParallelLinear(
            hidden_size,
            self.head_dim,
            self.total_num_heads,
            self.total_num_kv_heads,
            bias=qkv_bias,
        )
        self.o_proj = RowParallelLinear(
            self.total_num_heads * self.head_dim,
            hidden_size,
            bias=False,
        )
        self.rotary_emb = get_rope(
            self.head_dim,
            rotary_dim=self.head_dim,
            max_position=max_position,
            base=rope_theta,
            rope_scaling=None,
        )
        self.attn = Attention(
            self.num_heads,
            self.head_dim,
            self.scaling,
            self.num_kv_heads,
        )
        if not self.qkv_bias:
            self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
            self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        qkv = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
        q = q.view(-1, self.num_heads, self.head_dim)
        k = k.view(-1, self.num_kv_heads, self.head_dim)
        v = v.view(-1, self.num_kv_heads, self.head_dim)
        if not self.qkv_bias:
            q = self.q_norm(q)
            k = self.k_norm(k)
        q, k = self.rotary_emb(positions, q, k)
        o = self.attn(q, k, v)
        output = self.o_proj(o.flatten(1, -1))
        return output


class Qwen3MLP(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
    ) -> None:
        super().__init__()
        self.gate_up_proj = MergedColumnParallelLinear(
            hidden_size,
            [intermediate_size] * 2,
            bias=False,
        )
        self.down_proj = RowParallelLinear(
            intermediate_size,
            hidden_size,
            bias=False,
        )
        assert hidden_act == "silu"
        self.act_fn = SiluAndMul()

    def forward(self, x):
        gate_up = self.gate_up_proj(x)
        x = self.act_fn(gate_up)
        x = self.down_proj(x)
        return x


class Qwen3DecoderLayer(nn.Module):

    def __init__(
        self,
        config: Qwen3Config,
    ) -> None:
        super().__init__()
        self.self_attn = Qwen3Attention(
            hidden_size=config.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=config.num_key_value_heads,
            max_position=config.max_position_embeddings,
            rms_norm_eps=config.rms_norm_eps,
            qkv_bias=getattr(config, 'attention_bias', True),
            head_dim=getattr(config, 'head_dim', None),
            rope_theta=getattr(config, "rope_theta", 1000000),
            rope_scaling=getattr(config, "rope_scaling", None),
        )
        self.mlp = Qwen3MLP(
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
        )
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if residual is None:
            hidden_states, residual = self.input_layernorm(hidden_states), hidden_states
        else:
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
        hidden_states = self.self_attn(positions, hidden_states)
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual


class Qwen3Model(nn.Module):

    def __init__(
        self,
        config: Qwen3Config,
    ) -> None:
        super().__init__()
        self.embed_tokens = VocabParallelEmbedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList([Qwen3DecoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
    ) -> torch.Tensor:
        hidden_states = self.embed_tokens(input_ids)
        residual = None
        for layer in self.layers:
            hidden_states, residual = layer(positions, hidden_states, residual)
        hidden_states, _ = self.norm(hidden_states, residual)
        return hidden_states


class Qwen3ForCausalLM(nn.Module):
    packed_modules_mapping = {
        "q_proj": ("qkv_proj", "q"),
        "k_proj": ("qkv_proj", "k"),
        "v_proj": ("qkv_proj", "v"),
        "gate_proj": ("gate_up_proj", 0),
        "up_proj": ("gate_up_proj", 1),
    }

    def __init__(
        self,
        config: Qwen3Config
    ) -> None:
        super().__init__()
        self.model = Qwen3Model(config)
        self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size)
        if config.tie_word_embeddings:
            self.lm_head.weight.data = self.model.embed_tokens.weight.data

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
    ) -> torch.Tensor:
        return self.model(input_ids, positions)

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        return self.lm_head(hidden_states)

在整个推理流程中，模型接收到的输入是一个 token id 序列（图中的 token_ids），这些 token id 本身已经是 数字化表示，但它们还不能直接参与复杂的模型计算。因此，在进入模型内部之后，首先需要经过一个转换步骤，在图中我们可以看到：

token_ids → Embedding → hidden_state

也就是说，模型首先会通过 Embedding 层，将离散的 token id 转换为向量表示，这个向量表示就是我们在 Transformer 中经常看到的 hidden_state。

hidden_state 这个名字其实更多是一种 语义上的命名，之所以称为 hidden，是因为这些状态只存在于 模型内部的计算过程中，最终用户是无法直接看到它们的。

从用户角度来看，模型输入输出通常只有两种数据：

输入：token
输出：logits / token

而 hidden_state 则是模型内部在各层之间传递的中间表示。

从数据结构角度来看，hidden_state 本质上就是 一组高维向量（multi-dimensional tensor），例如 [sequence_length, hidden_size]，每一个 token 在模型内部都会被映射成一个向量表示，然后这些向量会在不同的 layer 之间不断更新。

那么，token id 是如何转换为 hidden_state 的呢？🤔

这个过程其实就是 Embedding，Embedding 的作用可以理解为 根据 token id，在词表中查找对应的向量表示，因此，这一步不仅需要 token id，还需要 vocabulary（词汇表） 来完成映射。

换句话说：token_id + vocab → embedding vector → hidden_state，完成这一步之后，数据就从离散的 token 表示，进入了模型内部统一使用的 向量表示空间。

在 Embedding 之后，数据就进入了模型中最核心的部分：Decoder Layers，在图中可以看到，Decoder Layer 会重复 num_layers 次，形成模型的主体结构。

每一层 Decoder Layer 内部主要包含两个核心模块：

Attention
MLP

这两个模块会不断地对 hidden_state 进行变换和更新。

与此同时，hidden_state 也会在这里真正与模型权重（weights）发生交互，也就是说 模型的 “能力” 其实主要体现在 Decoder Layer 的计算过程中，在这一阶段，大量的矩阵乘法、注意力计算等操作都会在 GPU 上执行。

当数据经过所有 Decoder Layers 之后，我们最终仍然得到的是 hidden_state，但这个向量表示仍然是模型内部的中间状态，用户是无法直接理解的。因此，在模型的最后，还需要经过一个 LM Head 模块，将 hidden_state 映射回 词表空间：hidden_state → LM Head → logits。

这里得到的 logits 表示的是 对词表中每一个 token 的概率评分。随后，推理系统会根据这些 logits 进行：

greedy decoding
sampling
beam search

最终生成下一个 token。

如果从整体角度来看，一个语言模型的推理流程其实可以概括为：

token_ids
   ↓
Embedding
   ↓
hidden_state
   ↓
Decoder Layer × N
   ↓
LM Head
   ↓
logits

其中：

Embedding 与 LM Head 主要负责 token 与向量之间的转换
Decoder Layers 则承担了几乎所有核心计算

因此，在接下来的内容中，我们不会过多展开 Embedding 与 LM Head 的细节，而是会重点关注 当数据已经变成 hidden_state 之后，Decoder Layer 在 GPU 上到底执行了哪些计算。

而在 Decoder Layer 内部，最重要的两个模块正是我们前面提到的：

Attention
MLP

接下来，我们就将进一步深入到 Layer 内部，逐步分析这两个核心组件的实现方式。

3. 深入 Layer 内部

在前一节中，我们已经从整体上理解了模型推理的 宏观流水线。接下来，我们将进一步深入到模型内部，从 Decoder Layer 开始，逐步分析模型在每一层中究竟做了哪些计算。

从图中可以看到，Decoder Layer 并不是单独的一层结构，而是一个重复堆叠的模块。在实际模型中，Decoder Layer 通常会被重复很多次，从而形成整个模型的主体结构。例如，在 nano-vllm 中使用的 Qwen 模型 中，就包含 24 个 Decoder Layer。

这些 Layer 在结构上是完全相同的，也就是说每一层执行的 计算流程是一致的，每一层包含的 模块结构也是一致的，唯一不同的是 每一层使用的权重参数是不同的。

换句话说，所有 Layer 的结构是共享的，但每一层都会对应模型权重文件中的 不同参数片段。因此，虽然每一层 Layer 执行的计算流程相同，但由于参数不同，它们在处理 token 表示时会产生 不同的变换效果。

从直观角度来看，可以把每一层 Layer 想象成对 token 表示进行 逐步加工和提炼的过程。每经过一层，token 的向量表示都会被进一步更新，从而逐渐形成更加丰富的语义表达。

为了帮助理解，我们可以用一个比较形象的类比，例如：第一层可能更关注基础的语言模式，后面的层逐渐捕获更复杂的语义关系，不过需要特别强调的是 这种类比只是帮助理解的方式，并不代表模型内部真的存在明确的功能分工。在实际的大语言模型中，每一层 Layer 所学习到的知识，都是在 训练过程中通过优化目标自然形成的，这些知识并不是由人类事先规划或设计好的。

因此，从目前的研究来看，我们很难准确地解释每一层 Layer 具体学到了什么知识、不同 Layer 在语义层面承担了什么功能，模型内部的这些表示结构，仍然是当前深度学习研究中的一个重要问题。但从工程角度来看，我们至少可以确定一点：模型的核心计算几乎全部发生在这些重复堆叠的 Decoder Layer 中。

因此，在接下来的内容中，我们将进一步深入到 Decoder Layer 的内部结构，重点分析其中两个最重要的计算模块：

Attention
MLP

它们也是 Transformer 架构中最核心的组成部分。

由于模型中的 Decoder Layer 在结构上是完全一致的，因此只需要理解 单个 Layer 的计算过程，就可以理解整个模型的工作方式。

下面我们来看第二张架构图：

对应的代码实现如下：

class Qwen3DecoderLayer(nn.Module):

    def __init__(
        self,
        config: Qwen3Config,
    ) -> None:
        super().__init__()
        self.self_attn = Qwen3Attention(
            hidden_size=config.hidden_size,
            num_heads=config.num_attention_heads,
            num_kv_heads=config.num_key_value_heads,
            max_position=config.max_position_embeddings,
            rms_norm_eps=config.rms_norm_eps,
            qkv_bias=getattr(config, 'attention_bias', True),
            head_dim=getattr(config, 'head_dim', None),
            rope_theta=getattr(config, "rope_theta", 1000000),
            rope_scaling=getattr(config, "rope_scaling", None),
        )
        self.mlp = Qwen3MLP(
            hidden_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            hidden_act=config.hidden_act,
        )
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        residual: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if residual is None:
            hidden_states, residual = self.input_layernorm(hidden_states), hidden_states
        else:
            hidden_states, residual = self.input_layernorm(hidden_states, residual)
        hidden_states = self.self_attn(positions, hidden_states)
        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
        hidden_states = self.mlp(hidden_states)
        return hidden_states, residual

在这个结构中，每一层 Decoder Layer 接收的输入都是一个 hidden_state，回顾前面的整体流水线 token_ids → Embedding → hidden_state，也就是说，token 在经过 Embedding 之后会被转换为向量表示 hidden_state，然后送入 第一个 Decoder Layer。

而每一个 Layer 在内部完成一系列计算之后，最终仍然会输出一个新的 hidden_state，因此整个模型实际上形成了一个循环结构 hidden_state → Layer → hidden_state → Layer → …，每一层都会对 hidden_state 进行一次新的变换，然后将结果传递给下一层继续处理。

在图中可以看到，hidden_state 的尺寸标记为 4096，这个数值并不是随意设置的，而是由 $hidden_size = num_heads × head_dim(dimension) \text{hidden\_size = num\_heads}\times \text{head\_dim(dimension)}$ 这个关系决定的，其中：

num_heads：注意力头的数量
head_dim：每个 head 的向量维度

因此 hidden_state 本质上是一个 高维向量表示。

例如在 Qwen 模型中 $num_heads = 32 , head_dim=128 \text{num\_heads}=32,\ \text{head\_dim=128}$ ，因此 $hidden_size = 32 × 128 = 4096 \text{hidden\_size} = 32 \times 128 = 4096$ ，这意味着每一个 token 在模型内部都会被表示为一个 4096 维向量，这个向量包含了模型对该 token 的各种语义特征表示。

当一个 4096 维的 hidden_state 进入 Decoder Layer 后，第一个阶段就是 Multi-Head Attention，在 Qwen 模型中 $num_heads = 32 , head_dim=128 \text{num\_heads}=32,\ \text{head\_dim=128}$ ，因此 Multi-Head Attention 会产生 32 个 attention head，每个 head 的输出维度为 128。

值得注意的是，很多人在看到 Multi-Head Attention 时，会误以为将 4096 维向量 直接拆分成 32 份，每份 128 维，但实际上 并不是这样，图中也特意标注了 Projection — Not Split，也就是说 Multi-Head Attention 的本质是 线性投影（projection），而不是简单的拆分。

在这一层中，每一个 head 都会拥有 自己的一组权重参数，计算过程可以理解为：

hidden_state × head_weight → head_output

换句话说，输入是 同一个 4096 维 token 表示，每个 head 使用 不同的权重矩阵，最终得到 不同的 128 维表示，因此可以得到 head_0 head_1 … head_31，每一个 head 都会生成一个 128 维向量，可以理解为每一个 head 都从不同角度对同一个 token 表示进行理解。

当所有 head 的计算完成后，会将这些结果重新拼接起来，于是我们再次得到一个新的 hidden_state（32 × 128 → 4096），但是这个新的 hidden_state 与原始输入已经不同了，原因是 每个 head 都对输入向量进行了不同的特征提取和变换，因此新的表示包含了 更加丰富的信息。

Multi-Head Attention 之所以称为 Attention（注意力），不仅仅是因为多头结构，更重要的是当前 token 会与 序列中之前的 token 表示进行交互计算，也就是说，模型不仅在处理当前 token，还会参考 token 1 token 2 … token t-1 这些历史的 token 表示，这也是为什么推理过程中需要维护 KV Cache。

为了更直观地理解 Attention，我们可以用一个简单的类比，把整个过程想象成一次 圆桌会议。当前 token 就像是 一个新的参与者，当它进入会议时，会先询问之前已经发言的参与者：他们提供了哪些信息、哪些信息更重要，然后再结合自己的信息，得到新的理解。

因此 Attention 的核心作用就是：让当前 token 能够动态地关注序列中其他 token 的信息。

经过 Multi-Head Attention 之后，我们最终得到一个新的 hidden_state，这个新的表示具有两个重要特点：

1. 每个 head 都对输入特征进行了新的变换

2. 表示中融合了 当前 token 与历史 token 的关系信息

因此，相比输入的 hidden_state，新的表示包含了 更加丰富的上下文信息。随后，这个新的 hidden_state 会继续进入 MLP 模块，完成下一阶段的计算。

第二个阶段就是 MLP，在介绍 MoE 之前，我们先来看一下传统稠密 MLP（Dense MLP）的做法。

所谓 Dense 指的是在整个 MLP 层中，所有 token 都会经过 同一个 MLP 模块 完成计算，也就是说这里不存在 expert 的路由与选择机制，而是统一走一条固定的计算路径。

与 Attention 相比，MLP 阶段的逻辑会更简单一些。因为在这一阶段，模型 不再与序列中其他 token 发生交互计算。从设计上看，模型认为在前面的 Attention 阶段 中，当前 token 已经充分吸收了来自其他 token 的上下文信息，因此到了 MLP 这里，重点就不再是 “和别人交互”，而是 “对自己当前的表示做进一步加工”。

换句话说，此时输入到 MLP 的 hidden_state 已经同时包含了当前 token 自身的信息和来自前序 token 的上下文信息。因此，MLP 的核心职责可以理解为：在当前 token 的内部表示上做进一步的非线性变换，提升表示质量。

从图中可以看到，进入 MLP 的 hidden_state 维度是 4096，而在 MLP 内部，它会先被映射到一个更高维的中间状态，例如这里的 11008，这个中间表示通常就叫做 intermediate_state。

也就是说，MLP 的一个典型过程可以概括为 4096 → 11008 → 4096，这里的 11008 也不是随意选择的，而是模型配置中预先定义好的一个超参数。那为什么要先从 4096 升到 11008，再重新压回 4096 呢？🤔

我们可以从一个比较直观的角度理解这个过程：4096 维的 hidden_state 虽然已经包含了很多信息，但它的表达能力仍然是有限的。把它升维到更高的中间空间，本质上是为了给模型提供一个 更大的表示空间，让它能够在这个更高维的空间中，对当前 token 的特征进行更加充分的组合、筛选和重构。

你可以把它类比成一张分辨率有限的图片，原始的 4096 维表示，就像是一张细节已经存在、但还不够清晰的图像；而升维到 11008 维，则相当于给模型提供了一个更宽裕的内部加工空间，让它可以把原本隐含的细节重新展开、重新组织，使表示更加丰富。

当然，这里的 “升维” 并不是凭空创造信息，而是通过当前层对应的 权重参数，对输入表示做一次更加复杂的非线性变换。也就是说，模型会利用这一层专属的参数，将输入的 hidden_state 投影到一个更大的中间空间中，在这个过程中重新组合已有信息，并突出更有价值的特征。

但 MLP 最终输出时，又必须重新回到 4096 维，这是因为在整个模型中，hidden_state 的维度是各个模块之间、各层之间约定好的统一规格。也就是说 Layer 接收的是 4096 维 hidden_state，Layer 输出的也必须是 4096 维 hidden_state，因此，MLP 必须在完成中间加工之后，再把结果投影回原来的隐藏维度。

那这时大家可能又有个疑问：既然最后又回到了 4096 维，那前面升到 11008 维是不是等于 “白做了”？🤔

其实并不是，关键在于：升维和降维并不是简单地 “放大再缩小”，而是伴随着权重参数和非线性激活共同完成的一次特征重组。模型在更高维空间中完成了特征增强之后，再压缩回 4096 维时，保留下来的通常是更有用、更有判别力的部分。

因此，这个过程更像是一种 特征提纯，先把表示展开到更大的空间，在更大的空间中完成更充分的组合与变换，再将其中更有价值的部分压缩回固定维度。

从这个角度来看，MLP 更像是在做一种 “自我加工” 或 “表示提纯”，它不像 Attention 那样负责吸收外部上下文，而是专注于提升当前 token 表示本身的质量。所以我们也可以把 Attention 和 MLP 的职责做一个很清晰的区分：

Attention：负责引入其他 token 的信息，建立上下文关联
MLP：负责在当前 token 表示内部做进一步加工，提升信息质量

这也说明了为什么模型权重会在不同阶段发挥不同作用：

在 Attention 阶段，权重主要帮助模型理解 “当前 token 与其他 token 之间的关系”
在 MLP 阶段，权重主要帮助模型提升 “当前 token 表示本身的质量”

这就是传统 Dense MLP 在 Decoder Layer 中所承担的核心作用。

那到目前为止，我们已经介绍了 传统稠密 Transformer Layer 的基本结构：Attention → MLP，在这种结构中，MLP 采用的是 Dense（稠密）结构，也就是说，所有 token 在进入 MLP 时都会经过 同一个 MLP 模块 完成计算。

随着模型规模不断扩大，MLP 的计算成本也会迅速增长，因此近年来一些新的大模型开始采用另一种结构即 MoE（Mixture of Experts），其大致结构如下图所示：

在传统 Dense MLP 中，计算流程是：

hidden_state (4096)
      ↓
intermediate_state (11008)
      ↓
hidden_state (4096)

而在 MoE 中，这个过程会被拆分成 多个专家（Expert）。

例如在上图的例子中，我们可以看到输入仍然是 hiddent_state (4096)，中间会经过一个 router，router 会选择若干 expert，每个 expert 负责一部分计算。与 Dense MLP 不同的是，这里的中间维度通常会变小，例如 2560，也就是说，每个 Expert 的 MLP 规模会更小。

MoE 中的 E（Experts） 指的就是多个专家模块，这些 Expert 本质上仍然是 MLP 结构，只不过现在不再只有一个 MLP，而是存在多个不同的 MLP，例如 expert_1 expert_2 … expert_8。

专家的数量是一个模型设计参数，例如可以设计 8 个专家，也可以设计 16 个专家，甚至可以设计 256 个专家。不同模型的选择会有所不同，例如 Mistral 在早期 MoE 模型中使用的专家数量相对较少，DeepSeek 等模型则使用了更多专家，具体采用多少专家，通常与训练策略和模型规模有关，并不存在一个固定的 “最佳答案”。

此外，如果每次计算时 所有 Expert 都参与计算，那么计算量其实与 Dense MLP 没有本质区别，因此 MoE 引入了另一个关键组件：Router（路由器）。Router 的作用是根据输入的 hidden_state，选择最合适的 Expert 来处理当前 token，例如 hidden_state → router → top-k experts。

在很多模型中，一个 token 通常只会激活 top-1 或 top-2 experts，而不是全部专家，这就是 MoE 的核心思想：只激活少数专家，而不是所有专家。

为了帮助大家更好的理解，我们可以用一个简单的类比，可以把这些 Expert 想象成不同领域的专家，例如体育专家、语言专家、数学专家、历史专家等。当一个问题进入系统时，router 会判断哪些专家最适合处理当前问题，例如，如果输入内容与体育相关，那么系统可能会优先激活 体育专家。

当然，在真实的模型中，并不存在人工定义的 “体育专家” 或 “数学专家”，每个 Expert 的能力并不是人类事先设定的，而是在 训练过程中自动学习形成的，这个类比只是帮助大家理解 为什么只激活部分 Expert 也可以完成计算。

当 Router 只激活少数 Expert 时，就会出现一个重要特性：并不是所有 Expert 都参与计算，因此 MoE 被称为 Sparse Model（稀疏模型），与之对应的是 Dense Model（稠密模型），两者的区别可以简单理解为：

Dense：所有参数都会参与计算
Sparse：只有部分参数会参与计算

因此在 MoE 中，大量 Expert 在一次推理过程中 并不会被激活。

MoE 最核心的优势在于：可以显著降低每个 token 的计算量，因为对于 Dense MLP 而言，每个 token 都要计算完整 MLP，而对于 MoE 每个 token 只需要计算少数 Expert，因此在相同计算成本下，可以构建 参数规模更大的模型。

例如一个 500B 参数的 Dense 模型在实际算力条件下可能难以训练，但如果使用 MoE 可以拥有 数百个 Expert，总参数量可以非常大，但每次只激活少数 Expert，从而在 可接受的计算成本 下训练更大的模型。

当然，MoE 也不是没有代价，如果在 相同参数规模 下进行比较，例如 Dense 70B vs MoE 70B，那么理论上 Dense 模型通常能够提供更稳定、更充分的信息表达，原因在于 Dense 模型的所有参数在每次计算中都会参与，而 MoE 中只有部分 Expert 会被激活。

因此 MoE 的真正优势并不是 “专家越多知识越丰富”，而是 在有限算力条件下，让模型规模能够继续扩展。我们可以把 MoE 理解为一种 工程上的权衡：牺牲部分计算完整性，换取更大的模型规模。

OK，那到这里，我们就完整理解了 一个 Decoder Layer 内部的计算结构。

一个 Layer 的输入是 hidden_state (4096)，内部会依次经过 Attention、MLP（Dense 或 MoE），最终输出仍然是 hidden_state (4096)。因此 Layer 之间可以形成一个非常自然的链式结构：

hidden_state
    ↓
Layer 1
    ↓
Layer 2
    ↓
...
    ↓
Layer N

每一层都会重复执行一次 注意力计算 和一次 表示增强（MLP），随着 Layer 不断堆叠，token 的表示会逐渐融合更多上下文信息，并不断被提纯和强化。

4. 与 KV Cache 数据面结合

OK，接下来我们来看另一个非常重要的优化机制：KV Cache。

首先我们需要回答一个问题：KV Cache 是在模型推理的哪个阶段发挥作用的呢？

从上图可以看到，在一个 Decoder Layer 中主要包含 Attention 和 MLP 两个部分，其中 MLP 阶段 无需缓存，Attention 阶段 才需要 KV Cache。

原因其实很简单，在 MLP 计算过程中，每一个 token 只是与当前的 模型权重 进行一次前向计算：hidden_state → linear → activation → linear，每一轮输入的 token 都是独立计算的，得到的结果也是一次性的，不会被后续 token 再次使用，因此 没有缓存价值。

而 Attention 阶段则完全不同，在自回归生成过程中：

第 2 个 token 需要计算 与 token 1 的注意力
第 3 个 token 需要计算 与 token 1、token 2 的注意力
第 4 个 token 需要计算 与 token 1、token 2、token 3 的注意力
…

如果每次都重新计算这些历史 token 的 Key / Value，就会产生大量 重复计算，因此推理引擎（如 vLLM）会将已经计算过的 Key / Value 缓存下来，这就是 KV Cache 的核心思想。

在上一篇文章中我们已经介绍过 Block Manager 的设计，如上图所示，Block Manager 主要运行在 CPU 侧，它负责：

管理 token 序列
分配 block
维护引用计数
做 prefix cache 复用

在 CPU 侧，block 中存储的仍然只是 token 的逻辑结构，例如：

block {hash: h_a, ref: 1}
[token, token, token]

但这些只是 CPU 数据结构，并不是 GPU 上真实存在的 KV Cache。

接下来我们要看的，是 GPU 显存中的 KV Cache 实际布局：

在这里插入图片描述

实现代码如下：

# nanovllm/layers/attention.py
import torch
from torch import nn
import triton
import triton.language as tl

from flash_attn import flash_attn_varlen_func, flash_attn_with_kvcache
from nanovllm.utils.context import get_context


@triton.jit
def store_kvcache_kernel(
    key_ptr,
    key_stride,
    value_ptr,
    value_stride,
    k_cache_ptr,
    v_cache_ptr,
    slot_mapping_ptr,
    D: tl.constexpr,
):
    idx = tl.program_id(0)
    slot = tl.load(slot_mapping_ptr + idx)
    if slot == -1: return
    key_offsets = idx * key_stride + tl.arange(0, D)
    value_offsets = idx * value_stride + tl.arange(0, D)
    key = tl.load(key_ptr + key_offsets)
    value = tl.load(value_ptr + value_offsets)
    cache_offsets = slot * D + tl.arange(0, D)
    tl.store(k_cache_ptr + cache_offsets, key)
    tl.store(v_cache_ptr + cache_offsets, value)


def store_kvcache(key: torch.Tensor, value: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, slot_mapping: torch.Tensor):
    N, num_heads, head_dim = key.shape
    D = num_heads * head_dim
    assert key.stride(-1) == 1 and value.stride(-1) == 1
    assert key.stride(1) == head_dim and value.stride(1) == head_dim
    assert k_cache.stride(1) == D and v_cache.stride(1) == D
    assert slot_mapping.numel() == N
    store_kvcache_kernel[(N,)](key, key.stride(0), value, value.stride(0), k_cache, v_cache, slot_mapping, D)


class Attention(nn.Module):

    def __init__(
        self,
        num_heads,
        head_dim,
        scale,
        num_kv_heads,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.scale = scale
        self.num_kv_heads = num_kv_heads
        self.k_cache = self.v_cache = torch.tensor([])

    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
        context = get_context()
        k_cache, v_cache = self.k_cache, self.v_cache
        if k_cache.numel() and v_cache.numel():
            store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)
        if context.is_prefill:
            if context.block_tables is not None:    # prefix cache
                k, v = k_cache, v_cache
            o = flash_attn_varlen_func(q, k, v,
                                       max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q,
                                       max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k,
                                       softmax_scale=self.scale, causal=True, block_table=context.block_tables)
        else:    # decode
            o = flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache,
                                        cache_seqlens=context.context_lens, block_table=context.block_tables, 
                                        softmax_scale=self.scale, causal=True)
        return o

上图中青色的部分就是我们在 KV Cache 的数据面显卡上面做的分配的逻辑，在 GPU 显存中，KV Cache 的存储方式可以理解为一个 三维结构：

[layer][K/V][block][token_vector]

它与 CPU Block Manager 的关系是 CPU block 中的 [token token token] 对应到 GPU 上时，会被展开为：

layer0
   K cache
   V cache

layer1
   K cache
   V cache

...

layer23
   K cache
   V cache

也就是说一个 CPU block 在 GPU KV Cache 中会被 复制多份：每个 layer 一份，每个 layer 包含两部分：K 和 V。原因是不同 layer 的 attention 计算结果是完全不同的，因此每一层都必须独立缓存。

例如，如果模型有 24 层 decoder layer，也就是说 KV Cache layers = 24，那么 GPU 上就会有 24 x K cache 和 24 x V cache，共 48 份缓存结构。

每个 block 在所有 layer 上都可以通过 block id 精确定位，这样在后续计算中如果某个 block 已经计算过就可以直接 读取 KV Cache，避免重复计算 attention。

需要注意的是，在 CPU Block Manager 中，一个 token 只是一个 离散 ID，例如 token_id = 15234，而在 KV Cache 中，它变成了 向量表示：

key  : [head_dim * num_heads]
value: [head_dim * num_heads]

例如 num_heads = 32，head_dim = 128，那么 vector_dim = 32x128 = 4096，这意味着 KV Cache 中每个 token 存储的是 4096 维向量，这也是 KV Cache 占用显存非常大的原因。

在 nano-vllm 的代码中，KV Cache 的写入并没有直接调用 CUDA API，而是使用 Triton Kernel 来完成 GPU 显存操作。

核心代码如下：

@triton.jit
def store_kvcache_kernel(
    key_ptr,
    key_stride,
    value_ptr,
    value_stride,
    k_cache_ptr,
    v_cache_ptr,
    slot_mapping_ptr,
    D: tl.constexpr,
):

这个 kernel 的核心逻辑其实非常简单：

1. 读取当前 token 的 key / value

key = tl.load(key_ptr + key_offsets)
value = tl.load(value_ptr + value_offsets)

2. 根据 slot_mapping 找到它在 KV Cache 中对应的 slot

slot = tl.load(slot_mapping_ptr + idx)

3. 将数据写入 KV Cache

tl.store(k_cache_ptr + cache_offsets, key)
tl.store(v_cache_ptr + cache_offsets, value)

整个过程就是 token → KV 向量 → 写入 KV Cache，而 slot_mapping 则由 Block Manager 提供，它负责将 token index → block slot 进行映射。

至此我们就把 KV Cache 的完整链路串起来了：

1. CPU 侧 Block Manager

管理 token block
分配 block id
维护引用计数

2. GPU 侧 KV Cache

按 layer 存储
每层包含 K / V 两部分
block 与 CPU block 一一对应

3. Triton Kernel

将 attention 计算得到的 K / V 向量
写入对应的 KV Cache slot

通过这种设计，推理引擎可以在生成过程中 复用历史 token 的 K/V，避免重复 attention 计算，从而大幅降低推理成本。

这也是 vLLM 能够实现 高吞吐推理的重要基础之一。

5. 实现 Tensor 并行

在理解了 KV Cache 的机制之后，我们再回到 每一个 decoder layer 内部的计算结构：

一个典型的 decoder layer 主要包含 Attention 和 MLP 两部分，在小模型中，这些计算通常可以在 单张 GPU 上完成。但当模型规模不断增大时，就会出现一个现实问题：单张 GPU 的显存不足以容纳全部模型权重。

例如：Attention 层需要加载 Q/K/V projection 权重，MLP 层需要加载 两层线性层权重，当模型参数规模达到几十亿甚至上百亿时，仅仅加载这些权重就可能超过单卡显存容量。

此时就需要一种机制，让 多张 GPU 协同完成一个 layer 的计算，这就是 Tensor Parallel（张量并行），Tensor 并行的核心思想是：将一个 layer 内部的计算拆分到多张 GPU 上执行，从而分摊权重存储和计算压力。

如上图所示，这里我们用 TP=2 的情况来说明 Tensor 并行的实现方式，在这个例子中 $num_heads = 32 , head_dim = 128 , hidden_state = 4096 \text{num\_heads}=32,\ \text{head\_dim}=128,\ \text{hidden\_state}=4096$ ，当 TP=2 时，我们会使用两张 GPU：GPU 0 和 GPU 1。

需要注意的第一点是 输入的 hidden_state 不会被拆分，也就是说 hidden_state (4096) 会被 完整复制到每一张 GPU 上，因此：

GPU0 input = hidden_state (4096)
GPU1 input = hidden_state (4096)

两张卡的输入是 完全相同的。

真正被拆分的其实是 Attention 的 projection 权重，在 TP=2 的情况下 32 heads → 16 + 16，于是 GPU0 只负责 head_0 ~ head_15，而GPU1 只负责 head_16 ~ head_31，这样每张 GPU 只需要加载 一半的 Attention 权重，这就使得 显存占用减少了一半。

由于每张 GPU 只拥有一半的 head 权重，因此它们只能计算 对应 head 的输出。例如在这个例子中，head_dim = 128，heads_per_gpu = 16，那么每张 GPU 计算得到的特征维度为 128 x 16 = 2048，但在实际实现中，输出仍然会被映射回 4096 维 hidden_state，只是此时得到的只是一个 partial output（部分输出）。

也就是说 GPU0 计算的 partial output 只包含 head_0 ~ head_15 的信息，GPU1 计算的 partial output 只包含 head_16 ~ head_31 的信息，虽然两个 GPU 的输出维度都是 4096，但它们包含的信息并不完整。

因此，在进入下一个 MLP 层之前，需要将两个 GPU 的计算结果进行合并，这个过程通常通过 All-Reduce（全聚合） 完成：

partial_output_gpu0 + partial_output_gpu1 = final_hidden_state

合并完成之后的 hidden_state (4096) 会再次被 同步到所有 GPU 上，这样每一张 GPU 都会得到 完整的 hidden_state，随后再进入下一层的计算。

这里有一个非常关键的问题：Tensor Parallel 的通信到底发生在哪一步呢？

我们可以分为两种情况来看：

1. 输入阶段（不需要通信）

当 hidden_state 输入到 Tensor Parallel 的 layer 时：

hidden_state → GPU0
hidden_state → GPU1

这里 不需要 GPU 之间通信，原因是每张 GPU 上都会持有 完整的 hidden_state，它只是使用不同的权重去进行计算，因此 输入阶段：无通信。

2. 输出合并阶段（需要通信）

当每张 GPU 计算出 partial output 之后：

GPU0 → partial_output_0
GPU1 → partial_output_1

此时每张 GPU 的输出都 只包含一部分信息，例如：GPU0：head_0 ~ head_15，GPU1：head_16 ~ head_31，为了得到完整的结果，我们需要执行 All-Reduce(partial_output) 操作。

这个操作会将所有 GPU 上的 partial output 进行求和，并将最终结果 同步回所有 GPU，最终得到完整的 hidden_state，此时每一张 GPU 都会拥有 一致的 hidden_state，可以继续进入下一层计算。

如果 Tensor Parallel 的规模进一步扩大，例如 TP=4、TP=8，那么通信模式也是类似的。例如 TP=4：

GPU0 partial
GPU1 partial
GPU2 partial
GPU3 partial

所有 GPU 会通过 All-Reduce 进行聚合：

partial_0 + partial_1 + partial_2 + partial_3

最终每张 GPU 都会得到 完整的 hidden_state，随着 GPU 数量增加 通信量会增加，通信延迟也会增加，因此 Tensor Parallel 在带来显存扩展能力的同时，也会引入 额外的通信开销。

Attention 的并行化讲完之后，我们接着看 MLP 的并行化，那相比 Attention，MLP 层的 Tensor Parallel 结构会更简单一些，MLP 的经典结构为：

hidden_state (4096)
      ↓
Linear (expand)
      ↓
activation
      ↓
Linear (project)
      ↓
hidden_state (4096)

在原始模型中 4096 → 11008 → 4096，当 TP=2 时第一层 Linear（expand）会进行 column parallel：

GPU0：4096 → 5504
GPU1：4096 → 5504

每张 GPU 只负责计算 一半的中间维度。

因此：

GPU0 → intermediate_state_0 (5504)
GPU1 → intermediate_state_1 (5504)

随后第二层 Linear（project）会进行 row parallel，此时每张 GPU 会得到一个 partial output，最后再通过 All-Reduce 将这些 partial output 合并，得到最终的 hidden_state (4096)。

通过 Tensor Parallel，我们可以让多张 GPU 协同完成一个 layer 的计算，这样带来的好处是每张 GPU 只需要存储部分权重 且多张 GPU 的显存可以 共同支撑更大的模型。

但与此同时，Tensor Parallel 也会带来一个新的成本：通信开销，这种通信主要发生在 All-Reduce 阶段，GPU 数量越多，通信开销也就越大。

因此在实际系统中，我们还需要在 显存扩展能力 和 通信开销 之间进行权衡。

OK，以上就是我们对于 Tensor 并行的一个理解了。

6. 思考

最后我们再回答两个比较有意思的问题：

第一个问题就是 模型中的 layer 数量和 attention head 数量分别代表什么呢？换句话说 什么时候需要更多 layer？什么时候需要更多 head？ 🤔

首先来看 layer（层数）。

在 Transformer 中，每一层基本都会重复同样的结构：Self-Attention → MLP，Self-Attention 的作用是 从上下文中提取和整合信息，而 MLP 的作用则是 对这些信息进行进一步的非线性加工。

因此，每增加一层 layer，本质上就相当于让模型 多进行一次信息整合与加工的过程，可以比较直观地理解为：layer 越多，模型能够进行的推理步骤就越多。

换句话说：layer ≈ 推理深度，更多的 layer 通常意味着更复杂的特征组合、更深层次的语义抽象以及更强的推理能力。但 layer 当然也不是越多越好，layer 增加会带来两个明显成本：第一个就是 计算量增加，每增加一层，推理计算量都会增加；第二个就是 收益递减，在实际模型设计中，人们发现当 layer 数量超过某个范围之后，性能提升会逐渐变慢。

因此当前很多模型的 layer 数量，其实是 在性能与计算成本之间的一种经验平衡。

接下来我们来看 attention head。

在 Multi-Head Attention 中，每一个 head 都会独立学习一组 Q/K/V 投影并计算一组注意力，可以理解为每个 head 都在关注 不同的语义关系。例如在自然语言中，不同 head 可能会关注语法关系、指代关系、长距离依赖
、关键词信息。

因此，多 head 的作用可以理解为：head ≈ 信息关注的不同视角，head 数量越多，模型就有机会从更多角度观察同一段文本并捕获更加丰富的语义关系。不过需要注意的一点是 head 数量的增加并不会直接增加 hidden size，例如： $hidden_size = num_heads × head_dim \text{hidden\_size} = \text{num\_heads} \times \text{head\_dim}$ ，如果 hidden_size 固定例如 $\times 128$ ，那么增加 head 的数量往往意味着每个 head 的维度会变小，因此 head 也同样存在一个 合理范围。

如果用一个比较直观的方式来理解，我们可以把模型结构看成两个维度：layer → 深度、head → 广度，layer 决定信息处理的深度，而 head 决定信息观察的多样性，但这并不意味着我们可以通过简单调整这两个参数，就能设计出某种 “风格” 的模型。

例如一个直觉上的想法可能是：如果我想做一个 “深度推理模型”，是不是可以减少 head，同时大幅增加 layer，比如设置 head = 4, layer = 128，这样做是否可以得到一个 专注深度推理的模型呢？

实际上，目前的一些研究表明：这种极端结构通常效果并不好，模型能力并不会因为简单地调整 layer/head 比例 而发生根本变化。

事实上，当前大模型的很多结构参数入 layer、head、expert 等等更多是 经验上比较稳定的配置，因为在大规模训练过程中 attention head 会自动学习不同功能，expert 也会自动形成不同 specialization，这些能力并不是人类提前明确设计出来的，而是 在大量数据训练过程中自然形成的结构分工。

某种程度上，这也与人类学习过程有些类似，一个人的认知能力往往不是来源于只学习一个非常狭窄的领域，而是来源于跨领域知识、触类旁通、举一反三，同样地，模型也需要在 深度和广度之间保持一定的平衡。

从目前主流的开源模型如 DeepSeek、Qwen、Mistral 来看，它们的 layer 与 head 的比例通常都不会非常极端，整体结构往往接近一种比较均衡的比例关系，例如 layer : head ≈ 1 : 1，即使存在差异，通常也不会超过 2 : 1，这说明模型结构并不是一个可以随意调节就改变模型能力的旋钮。

相比之下，目前真正影响模型性能的因素往往是 训练数据 和 训练方法，例如数据质量、数据规模、训练策略、对齐方法等，这些因素往往对模型性能产生更直接的影响。

因此，从当前实践来看：layer 决定模型的信息处理深度，head 决定模型关注信息的多样性，但模型能力的提升，并不能简单通过调整这两个结构参数实现。在现代大模型训练中，更重要的因素通常是 数据质量 + 训练方法，只要数据和训练方法足够好，并保持模型结构的 合理深度与广度，往往就能够得到性能优秀的模型。

第二个问题是 为什么近年来 MoE（Mixture of Experts）架构越来越流行了呢？🤔

其实原因非常简单，本质上还是一个 算力与参数规模之间的权衡问题。首先需要明确一点：MoE 的核心目标并不是直接提升模型质量，而是提升参数规模的可扩展性。

在传统的 稠密模型（Dense Model） 中，每个 token 都会激活所有参数，例如 70B dense model，那么每一次前向计算都需要使用 全部 700 亿参数，这会带来两个问题：计算量非常巨大、训练成本极高。

当模型规模继续扩大时，例如：100B、200B、500B，训练一个完全稠密的模型就会变得越来越困难。

MoE 的核心思想其实非常简单：每次只激活一部分专家（experts）进行计算，例如：

MoE
total params = 200B
active params per token = 20B

也就是说虽然模型总共有 200B 参数，但每个 token 只会使用其中的一小部分。这样就带来了一个非常关键的优势：参数规模 ↑ 计算量 ≈ 不变，因此 MoE 可以在 计算成本相近的情况下，训练出 更大规模的模型。

在推理场景下，我们可能会希望 直接使用一个非常大的稠密模型，以获得最好的输出质量，例如 dense 200B。但问题在于：这样的模型往往根本训练不出来，因为训练阶段需要消耗极其巨大的算力和显存资源，例如 GPU 数量、GPU 显存、通信带宽、训练时间等，如果基础设施不足，模型规模就会受到非常大的限制。

而 MoE 则提供了一种折中的方案：dense 200B → 难以训练，MoE 200B → 可以训练，虽然 MoE 在某些情况下 单次计算的表达能力不如完全稠密模型，但它能够让模型规模进一步扩大。而在当前阶段 更大的模型规模往往能够带来更好的性能，因此从工程角度来看 能训练出来 > 理论上更好但无法训练。

因此，从当前的发展趋势来看 MoE 的流行并不是因为它在理论上一定优于稠密模型，而是因为它在工程上提供了一种更好的 扩展路径，换句话说：

Dense Model → 计算效率高，但扩展困难
MoE Model → 稀疏计算，更容易扩展到超大规模

在当前算力仍然有限的情况下，MoE 架构成为了一种非常现实的选择。

OK，以上就是本次分享的全部内容了。

结语

在这一篇文章中，我们继续沿着上集的主线，把 nano-vllm 中原本作为 “黑盒” 的模型执行部分进一步展开，重点从模型结构和数据面的角度，分析了一次推理过程中真正发生的计算与缓存行为。

在开头部分，我们重新回到 Transformer 的 decoder layer 结构，对 attention 和 MLP 两个核心模块做了简单回顾，并进一步解释了为什么在推理过程中真正需要被缓存的是 attention 阶段的 Key / Value，而不是 MLP 的中间计算结果。也正是因为 attention 计算会不断依赖历史 token，因此 KV Cache 才成为推理引擎中最重要的一项性能优化。

随后我们从 Block Manager 的逻辑结构过渡到了 GPU 显存上的 KV Cache 数据布局。可以看到，在 CPU 侧 block 只是一个逻辑上的 token 容器，而在 GPU 上，每一个 block 都会被展开成按 layer 和 K / V 维度组织的缓存结构。通过这种方式，推理引擎可以在后续生成过程中直接复用历史 token 的 Key / Value，从而避免重复计算 attention，大幅降低推理开销。

接下来我们又进一步讨论了 Tensor Parallel 的实现逻辑。当模型规模超过单卡显存容量时，就需要通过多张 GPU 协同完成同一层的计算。通过对 attention head 和 MLP 中间维度进行切分，不同 GPU 可以分别计算各自的 partial output，并通过 all-reduce 等集合通信操作合并结果，从而在保证计算正确性的同时，实现模型权重在多卡之间的分布式存储与执行。

在最后的思考部分，我们简单讨论了两个经常被提到的问题：layer 数量和 attention head 数量分别意味着什么，以及为什么近年来 MoE 这样的稀疏模型架构越来越流行。可以看到，layer 更多代表模型的信息处理深度，而 head 则提供了更多不同视角的注意力表达；而 MoE 的流行，本质上则来自工程上的可扩展性—通过稀疏激活的方式，让模型能够在有限算力条件下继续扩大参数规模。

整体来看，下集的内容主要补齐了推理引擎中最底层的一块拼图：模型执行与 KV Cache 的数据组织方式。结合上集已经介绍过的 Scheduler、Block Manager、Model Runner 以及多卡并行系统，我们基本可以把一个 vLLM 风格推理引擎的核心链路完整地串联起来：从请求进入系统，到 sequence 的调度与批处理，再到 KV Cache 的管理、模型执行以及最终的生成输出。

nano-vllm 虽然只是一个相对精简的实现版本，但正因为它保留了 vLLM 中很多关键设计，同时又去掉了大量工程细节，因此非常适合作为理解现代 LLM 推理引擎内部结构的一个学习入口。希望通过这两篇文章的梳理，能够帮助大家对推理系统背后的核心思想建立一个更加清晰的整体认识🤗。

参考

加入AMD AI开发者计划！

免费领 200 小时云算力，进群参与显卡、AI PC 幸运抽奖

更多推荐

FlagOS Day 0 跨芯适配 GLM-5.2：发布即覆盖四款芯片，支持 vLLM + SGLang双插件

同时，FlagGems 新增 6 大领域算子库——FlagDNN、FlagBlas、FlagSparse、FlagFFT、FlagTensor、FlagAudio，覆盖科学计算与信号处理场景，共计 102 个领域算子，从"大模型专用"走向全领域覆盖。厂商目录放置后由插件自动发现加载，vLLM-Plugin-FL、SGLang-Plugin-FL、Megatron-LM-FL、Transformer