从零开始构建一个推理模型——使用预训练好的 LLM 生成文本

本章内容包括

• 为使用 LLM 搭建代码环境
• 如何使用 tokenizer 为 LLM 准备输入文本
• 使用预训练 LLM 进行文本生成的逐步过程
• 用于加速 LLM 文本生成的缓存与编译技术

在上一章中，我们讨论了传统大语言模型（LLM）与推理模型之间的区别，也介绍了多种提升 LLM 推理能力的技术。这些推理技术通常都是建立在传统（基础）LLM 之上的。

在本章中，我们将为后续章节打下基础，加载这样一个传统 LLM，后续我们会在它之上逐步应用各种推理技术，如图 2.1 所示。这个传统 LLM 已经完成了预训练，能够生成一般性文本，但它还没有被专门训练或增强为一个推理模型。

图 2.1 一个用于说明推理模型开发四个主要阶段的心智模型。本章聚焦于阶段 1：加载一个传统 LLM，并实现文本生成功能。

除了搭建编码环境和加载一个预训练 LLM，你还将学习如何使用 tokenizer 为模型准备文本输入。如图 2.1 所示，你还会实现一个文本生成函数，使这个 LLM 能够真正用于文本生成。这一功能会在后续章节中继续被使用并进一步改进。

2.1 面向文本生成的 LLM 简介

在本章中，我们会实现进行必要 LLM 实验所需的关键组件，搭建一个编码环境，并加载一个预训练 LLM，用它来生成文本，供后续章节使用。

像这样的 LLM 能够根据用户查询生成连贯的文本回应，如图 2.2 所示。

图 2.2 一个概览图，展示了 LLM 如何在接收用户查询（输入文本）后生成回答（输出文本）。

图 2.2 总结了一个 LLM 文本生成流水线中的关键组成部分，我们将在本章剩余部分讨论并实现这些步骤。

注意

按照惯例，关于神经网络的算法图通常是从左到右绘制的，因此数据会向前流动经过模型。

如果你已经在《从零开始构建大语言模型》里实现过这些内容，你会发现本章会很熟悉，因为我们会再次走一遍文本生成过程。不过，在本章里，我们会使用 PyTorch，而不是像那本书中那样只使用 NumPy，因为现在 PyTorch 已经成为这个领域更现实的默认选择。Sebastian Raschka 的《从零开始构建大语言模型》也对这些内容做了系统讲解。当然，理解那本书的细节并不是本章的前提，你仍然可以在本章中学到这些生成步骤的核心思想。

如果本章让你想起《从零开始构建大语言模型》中的相关内容，那是正常的，因为这里很多材料确实是为推理模型开发做准备的基础内容。

在真正开始实现图 2.2 中展示的组件之前，包括输入准备、加载预训练 LLM 以及生成文本，我们首先要确认编码环境已经就绪。这正是下一节的重点。

2.2 搭建编码环境

本节会介绍如何搭建代码环境，用于测试本书后续示例中的 Python 代码。我建议你在决定使用哪种方式之前，先把这一节完整读完。

如果你曾经写过 Python 代码，那么你很可能已经安装过 Python。在这种情况下，最简单的做法通常是确认你已经有一个 Python 环境可用（例如 Python 3.10 或更新版本），并且已经掌握 Python 的包安装器 pip 的基本使用。

如果你已经从本书出版社的网站下载了本书配套代码，那么可以通过 requirements.txt 一次性安装本章所需的 Python 依赖：

pip install -r requirements.txt

或者，你也可以不下载 requirements.txt 文件，而是直接执行：

pip install -r https://raw.githubusercontent.com/\
rasbt/reasoning-from-scratch/refs/heads/main/requirements.txt

本章使用的 Python 包

如果你更喜欢只安装本章用到的包，也可以直接运行下面这条命令：

pip install torch>=2.7.1 tokenizers>=0.21.2 reasoning-from-scratch

• torch 指的是 PyTorch，它是一个被广泛使用的深度学习库，提供了构建和训练神经网络所需的工具。
• tokenizers 是一个用于高效分词的库，我们会用它来为 LLM 准备输入文本。
• reasoning-from-scratch 是我为本书开发的一个自定义库。它包含贯穿全书的示例实现，以及我们会用到的一些辅助函数。

虽然 pip 依然是最经典、最常见的 Python 包安装方式，但我现在更推荐使用 uv。uv 是一个更现代的 Python 包和项目管理工具。简而言之，我更推荐它，因为它通常明显更快，而且它的依赖管理体验通常也比 pip 更稳定、更可靠。它还会自动创建隔离环境，并在某些情况下帮你处理系统 Python 与项目 Python 版本之间的兼容问题。它也非常适合新手，因为你可以在一个项目里快速开始，而不用先折腾一堆环境设置。

图 2.3 展示了一个四步过程，用于安装并使用 uv，以便运行本章中的代码，这一点我们会在本节后半部分详细展开。

图 2.3 通过 macOS 终端安装并使用 uv 这一 Python 包与项目管理工具。

图 2.3 所示流程是在 macOS 终端中完成的，不过 uv 同样支持 Linux 和 Windows。

1）安装 uv。你可以在 uv 的官方网站查看安装说明：
docs.astral.sh/uv/getting-…

2）接着，克隆 GitHub 仓库：

git clone --depth 1 https://github.com/rasbt/reasoning-from-scratch.git

这里的 --depth 1 选项会让 Git 执行一次浅克隆，这意味着它只下载该仓库的最新版本，而不会把完整的版本历史都拉下来。这样下载更快，也更省空间。

如果你的电脑上没有安装 Git，你也可以直接从出版社网站手动下载压缩包，或者在浏览器中访问下面这个地址下载 zip 文件：
github.com/rasbt/reaso…
（下载之后再解压。）

3）然后，在终端中进入 reasoning-from-scratch 文件夹。

4）在 reasoning-from-scratch 文件夹内部执行：

uv run jupyter lab

上面的命令会启动 JupyterLab，你也可以改成启动 Jupyter Notebook，然后执行本章对应的 notebook 代码。

提示

Python 脚本文件也可以通过下面这种方式执行：

uv run script-name.py

上面的 uv run ... 命令会自动调用一个虚拟环境（通常位于 .venv/ 文件夹里），并根据 reasoning-from-scratch 文件夹中的 pyproject.toml 自动安装所需依赖。所以，只要你用的是这个项目，你通常不需要额外手动安装 requirements.txt 里的依赖。

如果你还需要安装额外的包，可以使用如下命令：

uv add packagename

这本书配套仓库还包含一些额外的安装说明和注意事项，位于 ch02 文件夹下，供你在需要时参考。

2.3 理解硬件需求与建议

训练 LLM 非常昂贵。以当前主流 LLM 实践为例，训练一个模型通常需要数百万到上千万美元，而更大的模型甚至可能超过 5000 万美元的计算成本，才能在其基础上进一步叠加推理技术。

例如，Meta 的 Llama 3 模型——它本身也曾被不少推理工作当作基础模型——据报道是在 2,048 张 Nvidia H100 GPU 上训练约 11 周完成的，估算成本约为 550 万美元。（这里指的是完整训练成本，不包括后续进一步扩展和部署所需的额外成本。）

此外，根据相关技术报告，这类训练往往还会消耗数以千万计的 GPU 小时。如果你用一个非常粗略的平均值，比如 620 瓦功耗来估算，这相当于一个普通美国家庭大约 55 年的用电量。

这些资源需求在第一眼看上去可能非常吓人，尤其是对于想要学习如何构建 LLM 的开发者来说。但好消息是：我们并不需要重新训练一个完整的大模型。相反，我们可以从一个相对较小、已经预训练好的 LLM 出发，在其之上实现推理技术。

本章中我们使用的是一个较小的、基于 transformer 的解码器架构模型。后面章节介绍的方法，也同样适用于更大的 LLM。你可以把这里的过程理解成一种预算友好的实验路径。

总的来说，我的目标是：让你可以在现实可行的硬件条件下真正上手。就一般学习而言，你大概率不需要为了跑本章内容而专门买一台昂贵的新机器。相反，你甚至可以先从一个更小的测试环境开始，例如一台带有 Apple Silicon 芯片的新款 Mac mini，它依然足以让你学到很多关于推理与模型工作方式的核心内容。在我看来，先在一个较小系统里动手，往往更有助于真正理解这些机制，因为它会迫使你更认真地去观察计算与内存之间的权衡。

不过，尽管我们在本章里使用的是一个非常小的模型来做演示，但后续章节中的某些推理优化技术——尤其是第 5 到第 8 章中的一些内容——会更偏计算密集，因此在这些阶段，GPU 会更有帮助。

如果你跟着本节做环境测试，你应该已经安装好了 PyTorch。下面这段代码可以帮助你检查当前可用的计算设备类型：

import torch
 
print(f"PyTorch version {torch.__version__}")
if torch.cuda.is_available():
    print("CUDA GPU")
elif torch.xpu.is_available():
    print("Intel GPU")
elif torch.mps.is_available():
    print("Apple Silicon GPU")
else:
    print("Only CPU")

根据你的机器配置，代码可能返回：

PyTorch version 2.7.1
Only CPU

如果你是在一台只有 CPU 的机器上工作，那也完全没问题。第 2 到第 5 章的内容都可以在合理时间内运行完。

根据不同章节，代码会优先自动使用 NVIDIA CUDA GPU（如果可用），否则会退到 Apple Silicon GPU，再不然就使用 CPU。在后续各节和后续章节中，我也会提供更具体的说明。

如果你是研究人员，或者需要更强算力，但手头没有合适 GPU，也可以考虑云端选项。就我个人而言，我比较偏好 Lightning AI Studio（lightning.ai/），因为它用起来很顺手，功能也比较完整，如图 2.4 所示。另一个常见选择是 Google Colab（colab.research.google.com/）。

图 2.4 浏览器中的 Lightning AI GPU 云平台概览。该界面支持 Python 脚本、Jupyter notebook、终端访问，并允许用户根据计算需求在 CPU 与多种 GPU 之间切换。

截至写作本书时，Lightning AI 提供了一个用户友好的 Web 界面，你可以通过简单配置，在不同 GPU 选项之间灵活切换，如图 2.4 所示。（不过再次说明，本章并不要求 GPU；如果你确实想用 GPU，本章用一张 V4 GPU 就已经绰绰有余。）

注意

坦率地说，我曾在 2023 年帮助搭建并推出过 Lightning AI 平台上的一个功能，因此这里的推荐并非完全“零关联”。但我在这里推荐它，并不是因为我从中获得了任何回报，而是因为它确实是我个人实际在用、也确实觉得足够方便的一个选项。它支持快速切换 CPU 与 GPU，也支持临时环境和项目环境的快速启动。

如果你需要更详细的云平台使用建议，ch02 目录下还附带了一些额外说明。

关于 PyTorch

在本节中，我们导入并使用了 PyTorch。它是一套非常流行的通用神经网络库。全书中我们都会使用它，包括后面将实现的各种推理方法。

如果你对 PyTorch 还不熟悉，或者希望先快速掌握它的核心概念，我推荐你参考我网站上的课程材料《PyTorch in One Hour: Build and Train Neural Networks on Apple Silicon》（sebastianraschka.com/teaching/py…）。

2.4 为 LLM 准备输入文本

在本节中，我们会学习如何使用 tokenizer 来处理输入和输出文本，如图 2.5 所示。它展开展示了图 2.2 中更早提到的输入与输出准备步骤，从而更详细地说明 tokenizer 流水线的工作过程。

图 2.5 一个简化示意图，展示 LLM 如何接收输入数据并生成输出。用户提供的文本会先通过 tokenizer 的 encode 方法被转换为 token ID，然后送入 LLM。LLM 输出的 token ID 再通过 tokenizer 的 decode 方法还原成人类可读的文本。

为了让这一过程先在实践中跑通，我们会先从本书 reasoning-from-scratch Python 包里加载一个 tokenizer，这个包你应该已经在 2.2 节里安装好了。

我们先下载 tokenizer 文件（它对应的是 Qwen3 基础 LLM，后面几节我们会正式介绍它）：

from reasoning_from_scratch.qwen3 import download_qwen3_small
download_qwen3_small(kind="base", tokenizer_only=True, out_dir="qwen3")

这会显示一个类似下面的进度条：

tokenizer-base.json: 100% (6 MiB / 6 MiB)

这条命令会下载 tokenizer-base.json 文件（大约 6 MB），并将其保存到 qwen3 子目录下。

接下来，我们通过这个 tokenizer 文件来初始化一个 Qwen3Tokenizer：

from pathlib import Path
from reasoning_from_scratch.qwen3 import Qwen3Tokenizer
 
tokenizer_path = Path("qwen3") / "tokenizer-base.json"
tokenizer = Qwen3Tokenizer(tokenizer_file_path=tokenizer_path)

由于我们现在还没有真正加载 LLM（图 2.5 中展示的是完整组件流程），所以这里先只单独试用 tokenizer。具体来说，我们先做一个 tokenizer 的往返演示，也就是：先把一段文本编码成 token ID，再把这些 token ID 解码回字符串，如图 2.6 所示。

图 2.6 使用 tokenizer 演示“往返式”分词流程。用户提供的输入文本首先通过 encode 方法转换为 token ID，然后再通过 decode 方法准确地还原为原始文本。

下面两段代码分别实现图 2.6 中展示的编码和解码过程：

prompt = "Explain large language models."
input_token_ids_list = tokenizer.encode(prompt)

text = tokenizer.decode(input_token_ids_list)
print(text)

打印结果如下：

'Explain large language models.'

正如输出结果所示，编码再解码之后，tokenizer 成功重建出了原始输入提示词。

为了更直观地理解 tokenizer 是如何工作的，我们还可以逐个查看这些 token ID 分别对应的字符串：

for i in input_token_ids_list:
    print(f"{[i]} --> {tokenizer.decode([i])}")

输出如下：

840 --> Ex
20772 --> plain
3460 -->  large
4128 -->  language
4119 -->  models
13 --> .

如输出所示，原始文本会被拆成若干 token ID。这些 token 既可能表示整个词，也可能只表示词的一部分，这取决于 tokenizer 的切分方式。

例如，“Explain” 被拆成了两个 token：“Ex”和“plain”。这是因为 tokenizer 可能采用了一种基于 byte pair encoding（BPE） 的方法。BPE 会通过反复合并文本中最常见的字符组合，逐步构建词或子词单元。空格通常也会被包含进 token 中（例如 " large"），这样一来，LLM 在处理文本时就可以更自然地感知单词边界。

这个 Qwen3Tokenizer 拥有大约 151,000 个 token 的词表，这个规模相对相当大。作为对比，早期 GPT-2 的词表大约只有 50,000 个 token，而 Llama 3 的词表大约在 128,000 左右。

一个更大的词表，通常意味着语言模型在处理文本时能够用更少的 token 来表示更多内容，因为它更容易直接用完整词或更长的子词来表示文本，而不必频繁拆成很多很小的片段。比如，把 “Explain” 拆成 “Ex” 和 “plain” 这样的情况，在更大的词表下可能就不再需要。token 越少，输入序列通常越短，这会直接减少模型运行时的资源消耗。更具体地说，减少 token 数量往往会降低模型推理时的计算量，也会减少生成完整回答所需的 token 数量。

幸运的是，一个 tokenizer 的多 token 分词实现细节并不是本书的重点。不过，如果你对 tokenizer 机制本身很感兴趣，包括多字符实现等更深入内容，可以参考附录 C。

练习 2.1：编码未知词

试着用 tokenizer 自己做一些实验，看看它是如何处理未知词的。举例来说，如果你给它输入拼写错误的单词、虚构词语，或者不同语言里的单词，它会如何进行 token 化。

2.5 加载预训练模型

在上一节里，我们已经加载并初始化了 tokenizer，用它把输入文本转换成 token ID，并且又从 token ID 解码回文本。现在，在这一节里，我们将把真正的 LLM 模型也加载进来，如图 2.7 所示。

图 2.7 本书推理模型开发四个关键阶段的概览。本节聚焦于阶段 1：加载预训练 LLM。

正如上一节提到的，在本书中我们会使用 Qwen3 0.6B 作为一个轻量级预训练基础模型。在这一节中，我们就会真正把它加载进来，如图 2.7 所示。模型名里的 “0.6B” 表示：这个模型大约包含 6 亿个参数。

为什么选择 Qwen3？

虽然目前公开可用的开源模型有很多，我之所以选择 Qwen3 0.6B，主要有以下几点考虑：

• 对于本书来说，我们需要一个体量较小、适合在消费级硬件上运行的基础模型。
• 在同等参数规模下，Qwen3 模型家族在当前写作时通常具备较好的基准性能。
• Qwen3 0.6B 相比 Gemma 3 1B 和 Llama 2 1B 更节省内存。
• Qwen3 同时提供了基础模型版本（本书推理模型开发的重点）和指令模型版本，便于后续做参考与对比。

注意

这里所说的 “Qwen3” 指的是阿里巴巴的 Qwen3 系列模型，不是其他名字相似的模型。

为了强调我们是在“从零开始”搭建整个流程，我在本书的 GitHub 仓库中自己实现了一个轻量版 Qwen3 模型，而不是直接依赖外部现成的 LLM 库。这样做主要是出于教学目的：它可以让我们更清楚地看到模型内部是如何工作的，也更容易在后续实验中进行修改。尽管是从头实现的，但它依然与原始预训练 Qwen3 权重完全兼容。

不过，这并不意味着我们会在本章里深入拆解 Qwen3 的全部实现细节。这些内容会让本书偏离主线，更像是回到另一本书——《从零开始构建大语言模型》的范畴。本书的重点是：在一个基础模型之上实现推理方法，而在这里，这个基础模型就是 Qwen3。

注意

这个自行实现的 Qwen3 LLM 能完全在本地运行，这一点对很多读者来说也有额外价值，因为他们不必把自己的输入与输出发送给外部 API。如果你对隐私更敏感，那么本地运行会是一个非常有吸引力的选择。

如果你对 Qwen3 的更多背景信息感兴趣，可以查看附录 X。

在加载模型之前，我们还需要先决定：把模型放到哪个设备上运行，也就是 CPU 还是 GPU。清单 2.1 中的代码会自动选择当前可用的最佳设备：

清单 2.1 自动获取设备

def get_device(enable_tensor_cores=True):
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print("Using NVIDIA CUDA GPU")
        
        if enable_tensor_cores:
            major, minor = map(int, torch.__version__.split(".")[:2])
            if (major, minor) >= (2, 9):
                torch.backends.cuda.matmul.fp32_precision = "tf32"
                torch.backends.cudnn.conv.fp32_precision = "tf32"
            else:
                torch.backends.cuda.matmul.allow_tf32 = True
                torch.backends.cudnn.allow_tf32 = True
 
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        print("Using Apple Silicon GPU (MPS)")
 
    elif torch.xpu.is_available():
        device = torch.device("xpu")
        print("Using Intel GPU")
 
    else:
        device = torch.device("cpu")
        print("Using CPU")
 
    return device

如果你使用的是支持 Tensor Cores 的 NVIDIA GPU，并且使用的是较新的 PyTorch 版本，那么上面的 get_device() 函数会在 enable_tensor_cores=True 时自动启用 Tensor Cores，以便更快执行矩阵运算。这一设置通常只会带来极小的数值差异，但换来的速度提升往往是值得的。对于非 NVIDIA 设备，这些设置会被忽略。

接下来，我们可以直接调用它来拿到当前设备：

device = get_device()

很多硬件平台都支持一定程度的自动优化，因此即便你只是想调试代码或验证思路，让程序自动选择设备通常也是有帮助的。当然，为了便于复现实验和排查问题，你也可以显式指定设备：

device = torch.device("cpu")

不管用哪种方式，在继续本章之前，你都应该先确认模型确实被加载到了你想要的设备上。如果系统有 GPU，你通常会看到更好的性能。

注意

本章中的主要实验是在一台 Mac mini M4 上完成的；后文中提到的性能测试还包括 Apple Silicon M4 GPU 和 NVIDIA H100 GPU 的结果。

不过，在把模型放到设备上之前，我们首先要下载 Qwen3 0.6B 所需的权重文件。这些文件包含了初始化这一预训练模型所需的全部参数：

download_qwen3_small(kind="base", tokenizer_only=False, out_dir="qwen3")

输出如下：

qwen3-0.6B-base.pth: 100% (1433 MiB / 1433 MiB)
✓ qwen3/tokenizer-base.json already up-to-date

（这里 tokenizer 会显示一个勾选提示，因为我们在上一节里已经下载过它了。）

现在，权重下载完成后，我们可以实例化 Qwen3Model 类，然后通过 PyTorch 的 load_state_dict 方法把预训练权重加载进去：

from reasoning_from_scratch.qwen3 import Qwen3Model, QWEN_CONFIG_06_B
 
model_path = Path("qwen3") / "qwen3-0.6B-base.pth"
model = Qwen3Model(QWEN_CONFIG_06_B)
model.load_state_dict(torch.load(model_path))
model.to(device)

如果你的 device 设置是 "cpu"，那么 model.to(device) 这一操作可能会显得“多余”，因为默认情况下模型本来就在 CPU 内存里。但把这一步写上是个好习惯，因为后面切换设备时就不需要改代码。

执行完上面的代码后，你应该会得到类似如下输出（如果你是在交互式环境里运行，比如 Jupyter Notebook，那么只要打印 model 就可以看到）：

Qwen3Model(
  (tok_emb): Embedding(151936, 1024)
  (trf_blocks): ModuleList(
    (0-27): 28 x TransformerBlock(
      (att): GroupedQueryAttention(
        (W_query): Linear(in_features=1024, out_features=2048, bias=False)
        (W_key): Linear(in_features=1024, out_features=1024, bias=False)
        (W_value): Linear(in_features=1024, out_features=1024, bias=False)
        (out_proj): Linear(in_features=2048, out_features=1024, bias=False)
        (q_norm): RMSNorm()
        (k_norm): RMSNorm()
      )
      (ff): FeedForward(
        (fc1): Linear(in_features=1024, out_features=3072, bias=False)
        (fc2): Linear(in_features=1024, out_features=3072, bias=False)
        (fc3): Linear(in_features=3072, out_features=1024, bias=False)
      )
      (norm1): RMSNorm()
      (norm2): RMSNorm()
    )
  )
  (final_norm): RMSNorm()
  (out_head): Linear(in_features=1024, out_features=151936, bias=False)
)

这段输出是 PyTorch 打印出来的 Qwen3 0.6B 模型结构摘要。它突出了模型的关键组件：一个嵌入层、28 个 transformer block，以及一个最终输出投影层。每个 transformer block 中又包括 grouped-query attention 机制、一个多层前馈网络，以及贯穿其中的归一化层。

这些组件的整体结构在图 2.8 中做了直观展示，帮助读者更好地理解 LLM 架构。不过，本章并不要求你必须完全掌握这些结构细节。既然我们并不打算在这里修改模型本体，而只是要在它之上构建推理方法，你完全可以先把它当作一个黑盒来看待。当然，如果你对其中诸如 RMSNorm 之类的组件感兴趣，可以参考附录 T。

图 2.8 Qwen3 0.6B 模型架构概览。输入文本先经过 tokenizer，再进入嵌入层，然后通过 28 个重复的 transformer block。每个 block 内部都包含 grouped-query attention、前馈层与 RMS 归一化。模型最后以一个归一化层和一个线性输出层结束。图中的箭头展示了数据在模型中的流动路径。

到这里，本节的核心目标就已经完成了：我们已经成功加载了一个预训练模型，也大致了解了它的架构。理论上，它现在已经具备生成连贯文本的能力了。下一节中，我们将真正让它跑起来：把它用作一个基础的文本生成器。

2.6 理解顺序式 LLM 文本生成过程

既然我们已经加载好了一个预训练 LLM，下一步就是写一个函数，让这个 LLM 真的能够生成文本。这一点对后面实现各种推理增强方法非常关键，因此我们会在本章后续反复用到，如图 2.9 所示。

图 2.9 本书推理模型开发四个关键阶段的概览。本节会解释 LLM 中文本生成的核心概念，使我们能够实现一个文本生成函数，并在本章剩余部分真正使用这个预训练 LLM。

不过，在真正开始写这个文本生成函数之前——如图 2.9 所示，这是本章剩余部分的重点——我们先来理解 LLM 文本生成背后的基本机制。

所有现代 LLM 的文本生成，本质上都是一个顺序式过程：LLM 每次只生成一个新 token。这个过程通常被称为自回归生成（autoregressive generation），如图 2.10 所示。

图 2.10 LLM 中顺序式（自回归）文本生成的示意图。每一次迭代中，模型都会根据输入以及此前已经生成的 token 来预测下一个 token，然后把这个新 token 累积回输入中，继续生成连贯输出。

图 2.10 展示的是一个宽视角的自回归生成概览。图中每次只显示一个生成出的输出 token（图中叫 “new”），以及它如何逐步扩展输入提示。它的目的是帮助你直观理解基于 LLM 的文本生成究竟是如何一步一步往前走的。

聊天机器人界面

尽管图 2.10 展示的是底层逐 token 预测过程，但在一个真正的聊天界面里，实际机制会稍微复杂一些。模型依然是一次预测一个 token，但系统还会额外管理系统提示词、历史对话以及其他上下文，让模型在多轮对话中表现得更自然、更一致。这里我们聚焦的是单轮生成场景，也就是当前回答不依赖前面已有回答的情形。更复杂的多轮对话上下文管理，可以参考附录 D。

现在，为了把这个顺序生成过程看得更清楚，我们只看其中的一个迭代步骤，如图 2.11 所示。也就是说，我们先不考虑完整的连续生成，而只聚焦于单步预测。

图 2.11 对自回归文本生成过程中的单次迭代进行更细致的观察。LLM 会生成一个与输入序列相对应、但整体右移一位的输出序列。在每次迭代中，模型都会预测序列中的下一个 token。本质上，LLM 学会的是：一次一个 token 地继续补全输入提示。

为了进一步看清图 2.11 中这一单步生成过程，我们继续沿用 2.4 节里的示例提示词：

“Explain large language models.”

对应代码如下：

prompt = "Explain large language models."
input_token_ids_list = tokenizer.encode(prompt)
print(f"Number of input tokens: {len(input_token_ids_list)}")
 
input_tensor = torch.tensor(input_token_ids_list)
input_tensor_fmt = input_tensor.unsqueeze(0)
input_tensor_fmt = input_tensor_fmt.to(device)
 
output_tensor = model(input_tensor_fmt)
output_tensor_fmt = output_tensor.squeeze(0)
print(f"Formatted Output tensor shape: {output_tensor_fmt.shape}")

squeeze 和 unsqueeze：张量形状调整

张量是具有 n 个维度的数据结构。很多 PyTorch 函数和模型组件都要求输入张量具有特定维度，因此我们经常需要通过增减长度为 1 的维度来让张量形状匹配。

PyTorch 中的 .squeeze() 和 .unsqueeze() 操作，就专门用于改变张量形状：

• .unsqueeze() 会在指定位置新增一个长度为 1 的维度
• .squeeze() 会去掉长度为 1 的维度

举个例子，如果模型要求输入张量必须带有一个 batch 维度，而你的输入只是一个普通向量，就可以这样补一个维度：

example = torch.tensor([1, 2, 3]) 
print(example)
print(example.unsqueeze(0))

输出为：

tensor([1, 2, 3])
tensor([[1, 2, 3]])

这里 .unsqueeze(0) 在位置 0 新增了一个维度，把一个 1D 张量变成了形状为 (1, 3) 的 2D 张量。

反过来，.squeeze(0) 会移除第 0 个位置上长度为 1 的维度：

example = torch.tensor([[1, 2, 3]]) 
print(example)
print(example.squeeze(0))

输出为：

tensor([[1, 2, 3]])
tensor([1, 2, 3])

这在你需要把额外补上的那个维度去掉时会很有用。

上面那段代码执行后，输出如下：

Number of input tokens: 6
Formatted Output tensor shape: torch.Size([6, 151936])

正如我们看到的，我们给模型输入了 6 个 token，模型返回了一个 6 × 151,936 的矩阵。这个 6 对应的是输入中的 6 个 token；第二个维度 151,936，则对应模型的输出词表大小。

也就是说，输入序列中的每一个 token，都对应一个长度为 151,936 的向量。你可以把这个向量理解成：模型对整个词表中每一个候选 token 的“打分”。分数越高，模型就越倾向于把它当作下一个生成 token。

因此，为了得到模型要生成的下一个文本 token，我们就要从这个 6 × 151,936 的输出矩阵里，找到最后一个 token 所对应那一行中分数最高的位置，再把那个位置映射回 tokenizer 中的 token ID，如图 2.12 所示。

图 2.12 更细致地展示 LLM 在一次文本生成迭代中输出的原始分数，是如何被转换成一个 token ID，以及它所对应的文本表示的。

下面我们就用代码来完成图 2.12 所示的过程。

由于当前这个 LLM 是按照 next-token prediction 训练的，并且如图 2.11 所示，我们现在只关心最后一个位置对应的分数，因此可以通过 [-1] 索引拿到它：

last_token = output_tensor_fmt[-1].detach()
print(last_token)

这里 .detach() 的作用，是把这个张量从模型训练图中分离出来。更直白地说，就是：我们只是想查看模型最后一个 token 的输出分数，不需要再保留用于反向传播的额外信息。这样可以节省内存，也让后面的操作更轻量。

打印结果会是一组对应最后一个 token 的 151,936 个值：

tensor([ 7.3750,  2.0312,  8.0000,  ..., -2.5469, -2.5469, -2.5469],
       dtype=torch.bfloat16)

接下来，我们通过 argmax 找到最大分数的位置：

print(last_token.argmax(dim=-1, keepdim=True))

结果是：

tensor([20286])

这个整数值表示：在最后一个 token 对应的输出向量里，位置 20286 拿到了最高分。也就是说，模型预测下一个 token ID 应该是 20286。然后我们把它交给 tokenizer 去解码：

print(tokenizer.decode([20286]))

输出为：

Large

max 与 argmax

这里值得顺手回顾一下 max 和 argmax 的区别，因为我们后面会频繁用到 torch.argmax() 来选择下一个 token。

PyTorch 中的 torch.max() 和 torch.argmax() 功能不同：

• torch.max() 返回的是最大值本身
• torch.argmax() 返回的是最大值所在的位置索引

例如：

example = torch.tensor([-2, 1, 3, 1])
print(torch.max(example))
print(torch.argmax(example))

输出是：

tensor(3)
tensor(2)

最大值是 3，而它第一次出现的位置索引是 2。

如果你在 torch.argmax() 中使用 keepdim=True：

print(torch.argmax(example, keepdim=True))

则输出为：

tensor([2])

这里 keepdim=True 会保留输出张量的维度数，这在后续拼接张量、或者把 token ID 交回 tokenizer 的时候非常方便。

到目前为止：

• 图 2.10 展示了 LLM 的逐步生成整体过程；
• 图 2.11 放大了其中单次迭代；
• 图 2.12 又进一步说明：在这一轮迭代中，模型输出的分数矩阵是如何被转换成具体 token ID 的。

既然我们现在已经知道，LLM 是如何生成一个 token 的，下一节就可以把这些步骤串起来，写出一个真正能连续生成文本的最小文本生成函数。

2.7 编写一个最小文本生成函数

上一节解释了 LLM 中文本生成这一基本的顺序式、自回归生成过程。在本节中，我们将在这个概念基础上，真正实现一个文本生成函数，让这个预训练 LLM 能够根据给定提示生成连贯文本，如图 2.13 所示。

图 2.13 本书推理模型开发四个关键阶段的概览。本节会为预训练 LLM 实现一个文本生成函数。

图 2.13 中提到的这个文本生成函数，首先会把输入提示词转换为 token ID，交给模型处理。模型会预测最有可能的下一个 token，然后我们把它追加到当前序列中，再继续扩展这个序列来生成新的 token。这个迭代过程会一直持续，直到满足某个停止条件，然后我们再把生成出来的 token ID 解码成人类可读文本。

图 2.14 就展示了这个流程：它不仅显示生成出的 token ID，还同时显示它们对应的文本片段。

图 2.14 大语言模型（LLM）中顺序式（自回归）文本生成的示意图，并显式展示 token ID。每次迭代中，模型都会基于原始输入以及此前所有已生成 token 来预测下一个 token。这个预测出来的 token 会同时以文本形式和 token ID 形式被追加到序列中。

下面的 generate_text_basic_stream 函数，就是图 2.14 所描述的顺序生成过程的一个最基础实现。它使用上一节介绍过的 argmax 方式来决定每一步生成哪个 token。

清单 2.2 一个基础文本生成函数

@torch.inference_mode()
def generate_text_basic_stream(
    model,
    token_ids,
    max_new_tokens, 
    eos_token_id=None
):
 
    model.eval()
 
    for _ in range(max_new_tokens):
        out = model(token_ids)[:, -1]
        next_token = torch.argmax(out, dim=-1, keepdim=True)
 
        if (eos_token_id is not None
                and torch.all(next_token == eos_token_id)):
            break
 
        yield next_token
        
        token_ids = torch.cat([token_ids, next_token], dim=1)

清单 2.2 中的 generate_text_basic_stream 函数会通过 argmax 选出下一个 token，然后在一个固定的最大迭代次数（max_new_tokens）内不断重复这一过程。它返回的其实是一个个逐步生成的 token ID，就像图 2.14 中展示的那样，随后我们可以再把这些 token ID 解码成文本。

下面我们用这个函数，让 Qwen3Model 针对一个简单提示生成 100 个新 token：

prompt = "Explain large language models in a single sentence."
input_token_ids_tensor = torch.tensor(
    tokenizer.encode(prompt),
    device=device
    ).unsqueeze(0)
 
max_new_tokens = 100
 
for token in generate_text_basic_stream(
    model=model,
    token_ids=input_token_ids_tensor,
    max_new_tokens=max_new_tokens,
):
    token_id = token.squeeze(0).tolist()
    print(
        tokenizer.decode(token_id),
        end="",
        flush=True
    )

你可能需要等待 1 到 3 分钟不等，具体取决于你的计算机配置。

生成出来的文本如下：

Large language models are artificial intelligence systems that can
understand, generate, and process human language, enabling them to
perform a wide range of tasks, from answering questions to writing
articles, and even creating creative content.<|endoftext|>Human language
is a complex and dynamic system that has evolved over millions of
years to enable effective communication and social interaction. It is
composed of a vast array of symbols, including letters, numbers, and
words, which are used to convey meaning and express thoughts and
ideas. The evolution of language has

根据设备不同，你生成出的文本内容可能略有变化，例如在不同硬件上会出现一些细微差异。

从上面的输出可以看出，模型确实遵循了我们的指令，生成了一句关于“大语言模型”的单句描述。但它并没有在句号处停下来，而是继续生成，直到产生了一个特殊 token <|endoftext|>。这个 token 通常在训练过程中被用来表示一个文档或一个样本的结束。

提示

输出开头的 " Large"（第一个生成出的 token）之所以前面带空格，是因为 tokenizer 会把“前导空格”一起编码进 token 中。在清单 2.2 中，token 是一个接一个解码的，后面 token 会自然接在前面输出之后，因此最开始这个前导空格看上去会比较明显。如果你想让最终拼接出来的文本更自然，可以对第一个 token 用 .lstrip() 去掉左侧空格。

当我们在推理阶段使用模型时，通常希望它在产生 <|endoftext|> 这种特殊 token 后就停止。这个 token 对应的 ID 是 151643，我们可以确认一下：

print(tokenizer.encode("<|endoftext|>"))

更方便的是，这个值也已经保存在 tokenizer.eos_token_id 属性里。因此，我们只要把这个 ID 传给 generate_text_basic_stream，就能让生成在遇到结束标记时自动停止：

for token in generate_text_basic_stream(
    model=model,
    token_ids=input_token_ids_tensor,
    max_new_tokens=max_new_tokens,
    eos_token_id=tokenizer.eos_token_id
):
    token_id = token.squeeze(0).tolist()
    print(
        tokenizer.decode(token_id),
        end="",
        flush=True
    )

输出会变成：

Large language models are artificial intelligence systems that can
understand, generate, and process human language, enabling them to
perform a wide range of tasks, from answering questions to writing
articles, and even creating creative content.

这样一来，生成过程在检测到结束 token 后就会自动停住，不会再继续输出无关后文。

上面的简单示例已经能证明：我们的生成函数确实可以工作，并且能生成符合提示要求的文本。不过，它目前还非常慢，这一点你在自己的硬件上应该也会明显感觉到。

在真正讨论如何优化速度之前，我们先实现一个简单的辅助函数，用来统计文本生成过程中的时间和速度：

清单 2.3 token 生成速度与内存占用统计

import warnings
 
def generate_stats(output_token_ids, tokenizer, start_time,
                   end_time):
    total_time = end_time - start_time
    print(f"\n\nTime: {total_time:.2f} sec")
    print(f"{int(output_token_ids.numel() / total_time)} tokens/sec")
 
    for name, backend in (("CUDA", getattr(torch, "cuda", None)),
                          ("XPU", getattr(torch, "xpu", None))):
        if backend is not None and backend.is_available():
            device_type = output_token_ids.device.type
            if device_type != name.lower():
                warnings.warn(
                    f"{name} is available but tensors are on "
                    f"{device_type}. Memory stats may be 0."
                )
    
            if hasattr(backend, "synchronize"):
                backend.synchronize()
            
            max_mem_bytes = backend.max_memory_allocated()
            max_mem_gb = max_mem_bytes / (1024 ** 3)
            print(f"Max {name} memory allocated: {max_mem_gb:.2f} GB")
 
            backend.reset_peak_memory_stats()

generate_stats 函数会根据起止时间计算总耗时、平均每秒生成多少 token，以及 GPU 峰值显存占用。这里之所以同时检查 CUDA 和 XPU，是因为 PyTorch 对这两类后端提供了类似的工具函数，而 Apple Silicon GPU 和 MPS 目前没有完全相同的接口。

为了使用 generate_stats，我们需要在文本生成前后记录时间，并收集生成出的 token：

import time
 
start_time = time.time()
generated_ids = []
 
for token in generate_text_basic_stream(
    model=model,
    token_ids=input_token_ids_tensor,
    max_new_tokens=max_new_tokens,
    eos_token_id=tokenizer.eos_token_id
):
    token_id = token.squeeze(0).tolist()
    print(
        tokenizer.decode(token_id),
        end="",
        flush=True
    )
    next_token_id = token.squeeze(0)
    generated_ids.append(next_token_id)
 
end_time = time.time()
output_token_ids_tensor = torch.cat(generated_ids, dim=0)
generate_stats(output_token_ids_tensor, tokenizer, start_time, end_time)

在一台 Mac mini M4 CPU 上，输出如下：

Time: 7.94 sec
5 tokens/sec
 Large language models are artificial intelligence systems that can
understand, generate, and process human language, enabling them to
perform a wide range of tasks, from answering questions to writing
articles, and even creating creative content.

每秒 5 个 token 的生成速度显然非常慢。下一节中，我们会实现一种缓存技术，把这个生成过程提速大约 5 到 6 倍。

文本生成与推理（inference）术语

在 LLM 文献中，作者通常会自动使用术语 inference 来指代文本生成。这里的 “inference” 指的是：模型在已经学好参数之后，执行一次前向计算，并据此给出预测。

在神经网络语境中，inference 更具体地说，就是：拿一个已经训练好的模型，用它执行一次前向传播，然后输出预测结果，比如生成下一个 token。这和训练阶段是相对的；训练是在更新参数，而 inference 是用已经固定好的参数做推断。

如果换到统计学语境里，inference 还有更一般的含义：指的是从已知信息中推断未知信息，例如估计参数、分析不确定性，或者提出关于总体的假设。

但在大语言模型的上下文中，我们说的 inference，通常就是在给定输入后，让模型输出下一个 token 的概率分布，然后从中选出或采样一个 token。

2.8 通过 KV 缓存加速推理

现在我们已经实现了一个基础文本生成函数，也亲眼看到了它为什么慢。这就自然引出了一个非常重要的话题：推理性能优化。

在运行 LLM 推理时，也就是让模型根据提示生成文本时，性能很快就会成为瓶颈。很多现实系统都会采用工程技巧来让推理变得更高效。本节中，我们会介绍其中一个最基础也最有效的技巧：KV caching，如图 2.15 所示。

图 2.15 本书推理模型开发四个关键阶段的概览。本节建立在已加载好的预训练 LLM 和前面实现的基础文本生成函数之上，利用 KV 缓存来加速执行。

如图 2.15 所示，KV 缓存是本章介绍的首个推理优化技巧。这里的 KV 指的是 attention 机制中的 keys 和 values。如果你对这些术语还不太熟悉，不用担心；你现在只需要先理解核心思想：我们可以把某些中间结果缓存下来，在后续 token 生成时直接复用，从而加快整个生成过程，如图 2.16 所示。

图 2.16 KV 缓存在自回归文本生成中提升效率的示意图。与其在每一步都重新处理整个输入序列，不如把中间表示缓存下来，这样 LLM 在生成下一个 token 时就可以直接复用它们。这也避免了在后续每次迭代中都把新 token 与旧输入重新拼接后再整体送进模型。

使用 KV 缓存的核心思想——如图 2.16 所示——就是：在第一次前向传播时，把 attention 所需的一部分中间结果保存在缓存里。此前，每新生成一个 token，我们都需要把它与整个已有输入重新拼接，再把完整序列重新送入模型，这会带来大量重复计算。而使用 KV cache 之后，我们就能避免这种重复，只需要把新 token 输入模型即可，之前已有上下文所对应的中间 attention 表示则直接从缓存里取出来复用。

正如前面提到的，标准 decoder-only LLM 通常都会支持 KV 缓存，我们这里也会把它加进来，用来提升 token 生成速度。不过，如果你想更深入理解 KV cache 的具体实现细节，我还写过一篇专门文章：
Understanding and Coding the KV Cache in LLMs from Scratch
magazine.sebastianraschka.com/p/coding-th…

下面给出的是一个带 KV cache 的 generate_text_basic_stream 改造版本。它与清单 2.2 的基础版非常相似，只是增加了缓存逻辑。

清单 2.4 带 KV 缓存的基础文本生成函数

from reasoning_from_scratch.qwen3 import KVCache
 
@torch.inference_mode()
def generate_text_basic_stream_cache(
    model,
    token_ids,
    max_new_tokens,
    eos_token_id=None
):
 
    model.eval()
    cache = KVCache(n_layers=model.cfg["n_layers"])
    model.reset_kv_cache()
 
    out = model(token_ids, cache=cache)[:, -1]
    for _ in range(max_new_tokens):
        next_token = torch.argmax(out, dim=-1, keepdim=True)
 
        if (eos_token_id is not None
                and torch.all(next_token == eos_token_id)):
            break
 
        yield next_token
        out = model(next_token, cache=cache)[:, -1]

这个版本与清单 2.2 的主要区别在于：它引入了一个 KVCache 对象。

在第一次迭代里，模型仍然需要像以前那样，完整处理整段输入序列，即：

model(token_ids, cache=cache)

此时，模型会在计算过程中把输入序列对应的中间 KV 表示存进 cache 里。

而在后续迭代中，我们就不再需要每次都重新把完整序列送进去。相反，我们只需要把刚刚生成出来的 next_token 单独送入模型：

model(next_token, cache=cache)

模型会自动从 KV cache 中读取此前上下文的表示，并只为这个新 token 计算新增部分。

我们也可以像前一节那样，用同样的方式测试它的速度：

start_time = time.time()
generated_ids = []
 
for token in generate_text_basic_stream_cache(
    model=model,
    token_ids=input_token_ids_tensor,
    max_new_tokens=max_new_tokens,
    eos_token_id=tokenizer.eos_token_id
):
    token_id = token.squeeze(0).tolist()
    print(
        tokenizer.decode(token_id),
        end="",
        flush=True
    )
    next_token_id = token.squeeze(0)
    generated_ids.append(next_token_id)
 
end_time = time.time()
 
output_token_ids_tensor = torch.cat(generated_ids, dim=0)
generate_stats(output_token_ids_tensor, tokenizer, start_time, end_time)

输出如下：

Time: 1.40 sec
29 tokens/sec
 
 Large language models are artificial intelligence systems that can
understand, generate, and process human language, enabling them to
perform a wide range of tasks, from answering questions to writing
articles, and even creating creative content.

正如我们看到的，这个版本明显快了很多：在一台 Mac mini M4 CPU 上，它从原来的大约 5 tokens/sec 提升到了 29 tokens/sec。

更重要的是，生成出来的文本内容与之前完全一致，这并不奇怪，因为 KV cache 只是复用了中间结果，并没有改动模型本身的预测逻辑。

在下一节中，我们还会继续介绍另一种进一步加速的方法。它可以和这里的 KV cache 配合使用，也可以单独使用，并且对后续更复杂的推理方法也会很有帮助。

2.9 通过 PyTorch 模型编译进一步加速推理

在上一节中，我们已经介绍了 KV caching 这一加速技巧。现在，我们再来看另一种提升推理效率的方法：模型编译，如图 2.17 所示。

图 2.17 本书推理模型开发四个关键阶段的概览。本节建立在预训练 LLM 和前面实现的基础文本生成函数（包括 KV cache）之上，进一步加入模型编译以提升执行速度。

如图 2.17 所示，在本章接近尾声的时候，我们再引入一种通常能够进一步提速的技术：使用 torch.compile 对模型进行编译。这个特性允许模型在运行前被提前编译，从而减少运行时开销，并提升推理阶段的整体性能。

major, minor = map(int, torch.__version__.split(".")[:2])
if (major, minor) >= (2, 8):
    # This avoids retriggering model recompilations 
    # in PyTorch 2.8 and newer
    # if the model contains code like self.pos = self.pos + 1
    torch._dynamo.config.allow_unspec_int_on_nn_module = True
 
model_compiled = torch.compile(model)

如果你在 Apple Silicon 上运行时遇到 InductorError，请跳到第 2.9 节后面那部分说明。

值得注意的是：第一次执行编译后的模型，通常会更慢，因为系统还需要完成初始编译和优化步骤。为了更公平地衡量性能提升，我们会让文本生成过程运行多次。

先来看不带 KV cache 的版本。下面这段代码与前面类似，只是我们这次连续运行它三次，第一轮作为预热：

清单 2.5 使用编译模型进行文本生成

for i in range(3):
 
    start_time = time.time()
 
    generated_ids = []
    
    for token in generate_text_basic_stream(
        model=model_compiled,
        token_ids=input_token_ids_tensor,
        max_new_tokens=max_new_tokens,
        eos_token_id=tokenizer.eos_token_id
 
    ):
        token_id = token.squeeze(0).tolist()
        print(
            tokenizer.decode(token_id),
            end="",
            flush=True
        )
    
        next_token_id = token.squeeze(0)
        generated_ids.append(next_token_id)
    
    end_time = time.time()
    
    if i == 0:
        print("\n\nWarm-up run")
    else:
        print(f"\n\nTimed run {i}:")
 
    output_token_ids_tensor = torch.cat(generated_ids, dim=0)
    generate_stats(output_token_ids_tensor, tokenizer, start_time, end_time)
 
    print(f"\n{30*'-'}\n")

输出如下：

Large language models are artificial intelligence systems that can
understand, generate, and process human language, enabling them to
perform a wide range of tasks, from answering questions to writing
articles, and even creating creative content.
 
Warm-up run
 
Time: 11.68 sec
3 tokens/sec
 
------------------------------
 
 Large language models are artificial intelligence systems that can
understand, generate, and process human language, enabling them to
perform a wide range of tasks, from answering questions to writing
articles, and even creating creative content.
 
Timed run 1:
 
Time: 6.78 sec
6 tokens/sec
 
------------------------------
 
 Large language models are artificial intelligence systems that can
understand, generate, and process human language, enabling them to
perform a wide range of tasks, from answering questions to writing
articles, and even creating creative content.
 
Timed run 2:
 
Time: 6.80 sec
6 tokens/sec
 
------------------------------

如上所示，编译后的模型速度略有提升：从原来大约 5 tokens/sec，提升到大约 6 tokens/sec。

接下来，再看编译模型与 KV cache 一起使用时的表现。下面这段代码和前面基本一样，只是把 generate_text_basic_stream 换成了 generate_text_basic_stream_cache：

清单 2.6 使用 KV 缓存的编译模型进行文本生成

for i in range(3):
    start_time = time.time()
    generated_ids = []
    
    for token in generate_text_basic_stream_cache(
        model=model_compiled,
        token_ids=input_token_ids_tensor,
        max_new_tokens=max_new_tokens,
        eos_token_id=tokenizer.eos_token_id
 
    ):
        token_id = token.squeeze(0).tolist()
        print(
            tokenizer.decode(token_id),
            end="",
            flush=True
        )
    
        next_token_id = token.squeeze(0)
        generated_ids.append(next_token_id)
    
    end_time = time.time()
 
    if i == 0:
        print("\n\nWarm-up run")
    else:
        print(f"\n\nTimed run {i}:")
 
    output_token_ids_tensor = torch.cat(generated_ids, dim=0)
    generate_stats(
        output_token_ids_tensor, tokenizer, start_time, end_time
    )
 
    print(f"\n{30*'-'}\n")

输出如下：

Large language models are artificial intelligence systems that can
understand, generate, and process human language, enabling them to
perform a wide range of tasks, from answering questions to writing
articles, and even creating creative content.
 
Warm-up run
 
Time: 8.07 sec
5 tokens/sec
 
------------------------------
 
 Large language models are artificial intelligence systems that can
understand, generate, and process human language, enabling them to
perform a wide range of tasks, from answering questions to writing
articles, and even creating creative content.
 
Timed run 1:
 
Time: 0.60 sec
68 tokens/sec
 
------------------------------
 
 Large language models are artificial intelligence systems that can
understand, generate, and process human language, enabling them to
perform a wide range of tasks, from answering questions to writing
articles, and even creating creative content.
 
Timed run 2:
 
Time: 0.60 sec
68 tokens/sec
 
------------------------------

从上面的结果可以看到：在这台 Mac mini M4 上，使用 编译模型 + KV 缓存 后，文本生成速度从原来的 29 tokens/sec 一下提升到了 68 tokens/sec，几乎是又快了一个 2 倍多。

如果你的机器上效果没这么明显，可以尝试在 torch.compile 中使用 "max-autotune" 模式，而不是默认模式。例如，把：

model = torch.compile(model)

替换为：

model = torch.compile(model, mode="max-autotune")

练习 2.2：在非 CPU 设备上重新运行代码

如果你有 GPU，请尝试把本章的代码在 GPU 设备上重新运行，并比较 CPU 与 GPU 上的耗时差异。

如果你对性能更感兴趣，也可以把不同生成配置（普通版、KV cache 版、编译版）在不同设备上的表现做成一张表，类似于表 2.1。

表 2.1 不同硬件与模型配置下的 token 生成速度与显存占用

模式	硬件	tokens/sec	GPU memory
Regular	Mac Mini M4 CPU	5	-
Regular compiled	Mac Mini M4 CPU	6	-
KV cache	Mac Mini M4 CPU	28	-
KV cache compiled	Mac Mini M4 CPU	68	-
Regular	Mac Mini M4 GPU	27	-
Regular compiled	Mac Mini M4 GPU	43	-
KV cache	Mac Mini M4 GPU	41	-
KV cache compiled	Mac Mini M4 GPU	71	-
Regular	NVIDIA H100 GPU	51	1.55 GB
Regular compiled	NVIDIA H100 GPU	164	1.81 GB
KV cache	NVIDIA H100 GPU	48	1.52 GB
KV cache compiled	NVIDIA H100 GPU	141	1.81 GB

如上表所示，编译版通常会带来不错的性能提升，这一点是符合预期的。不过，关于 KV cache，也有一个值得特别说明的细节：在本章这种非常小模型 + 很短上下文的设置下，它的表现有时会看起来有些反直觉。

首先，本章使用的模型本来就很小。更大的模型通常更能从 KV cache 和编译带来的内存布局优化中受益，因为 GPU 可以更充分地发挥并行计算能力。也就是说，在更大的模型、更长的上下文、以及更重的推理负载下，这些优化往往会体现得更加明显。

其次，本例中的输入只有一个 prompt（即 batch size 为 1）。不同批量大小下的性能表现会有所不同，多输入批处理场景下的推理讨论可以参考附录 F。

2.10 小结

使用 LLM 生成文本，通常包括以下几个关键步骤：

• 搭建代码环境，以便运行 LLM 代码并安装必要依赖
• 加载一个预训练基础 LLM（例如 Qwen3 0.6B），后续我们会在它之上扩展推理能力
• 初始化并使用 tokenizer，将文本输入转换成 token ID，再把输出解码回人类可读文本
• LLM 的文本生成遵循顺序式（自回归）过程：模型一次生成一个 token，通过预测最可能的下一个 token 来逐步扩展序列

文本生成的速度和效率，可以通过以下方式提升：

• KV caching：缓存中间状态，避免在每一步都重新计算此前已经出现过的输入 token
• 模型编译（torch.compile） ：对运行时性能进行优化

本章通过实现一个可运行、较高效的基础文本生成流水线，为后续章节中的推理能力开发打下了技术基础。这个流水线建立在一个预训练基础 LLM 之上，而后续所有推理技术，都会在它之上继续展开。