构建推理模型——使用预训练 LLM 生成文本This chapter covers （本章内容）为使用 LLM 搭建代

This chapter covers

（本章内容）

为使用 LLM 搭建代码环境
使用 tokenizer 为 LLM 准备输入文本
通过预训练 LLM逐步实现文本生成的全过程
用缓存与编译技巧加速 LLM 文本生成

在上一章中，我们讨论了常规大型语言模型（LLM）与推理模型的差异，并引入了若干提升 LLM 推理能力的技术。这些推理技术通常是叠加在一个常规（基础）LLM之上的。
本章将为后续章节打下基础：我们会加载这样一个常规 LLM，后续章节会在其之上应用推理技术（见图 1）。这个常规 LLM 已经过预训练、能够生成通用文本，但尚未针对“推理”进行专门训练或增强。

Figure 2.1 A mental model depicting the four main stages of developing a reasoning model. This chapter focuses on stage 1, loading a conventional LLM and implementing the text generation functionality.
（用于构建推理模型的四个主要阶段的心智模型。本章聚焦阶段 1：加载常规 LLM，并实现文本生成功能。）

除搭建编码环境与加载预训练 LLM 外，你还将学习如何使用tokenizer为模型准备输入文本。正如图 2.1 所示，你还会实现一个文本生成函数，以便实际驱动 LLM 生成文本；该功能将在后续章节继续使用并加以改进。

2.1 Introduction to LLMs for text generation

（文本生成用 LLM 简介）

本章我们会实现 LLM 的所有基础要件：从搭建编码环境、加载预训练 LLM，到实现文本生成；这些组件会在本书中反复复用与扩展。从这个意义上说，本章也可以看作一个环境准备章节。
该 LLM 能够遵循基本指令并生成连贯文本（见图 2.2）。

Figure 2.2 An overview depicting an LLM generating a response (output text) given a user query (input text)
（概览：LLM 接收用户查询〔输入文本〕并生成响应〔输出文本〕。）

图 2.2 概括了一个 LLM 文本生成流水线的组成部分，我们将在本章稍后详细讨论并实现这些步骤。

NOTE
按惯例，涉及神经网络（如 LLM）的示意图通常自下而上阅读：底部为输入，顶部为输出；箭头表示信息在模型中的向上流动。
如果你此前没有以编程方式使用过 LLM，本章会讲清楚文本生成流程；不过，本章不会深入 LLM 的内部机理（如 attention 机制与其它结构组件）——这些内容属于我另一本书 Build a Large Language Model (from Scratch) 的主题。理解这些内部机制并非阅读本书的先决条件；如果你感兴趣，可以在完成本书后再深入学习。

在实现图 2.2 中的组件（输入准备、加载 LLM、生成文本）之前，我们先来搭建编码环境——这就是下一节的重点。

2.2 Setting up the coding environment

（搭建编码环境）

本节给出为本书示例准备 Python 环境的说明与建议。建议先完整阅读本节，再决定采用哪种方式。
如果你已经使用过 Python，且本地已有 Python 3.10+ 的环境，最简单的方式是在终端使用 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 PACKAGES USED IN THIS CHAPTER
（本章用到的 Python 包）

如果你只想安装本章用到的包，可执行：

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

torch：即 PyTorch，常用的深度学习库，提供构建与训练神经网络的工具。
tokenizers：高效分词库，用于为 LLM 准备输入数据。
reasoning-from-scratch：本书配套的自定义库，包含全书实现的代码示例与我们将使用的辅助函数。

虽然 pip 是安装 Python 包的标准方式，但我更推荐使用广受好评的 uv（Python 包与项目管理器）。uv 自带 Python 可执行文件，对尚未安装 Python 的系统也很友好。
图 2.3 概述了从安装 uv 到就绪运行本章代码的四步流程；本节接下来的内容将逐步讲解。

Figure 2.3 Installing and using the uv Python package and project manager via the macOS terminal
（在 macOS 终端安装并使用 uv。）

注意：图 2.3 演示的是在 macOS 终端中的安装与使用；uv 同样支持 Linux 与 Windows。

按官方文档为你的操作系统安装 uv：
docs.astral.sh/uv/getting-…
克隆 GitHub 仓库：

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

其中 --depth 1 表示浅克隆：只下载最新版本代码而不含完整历史，速度更快且占用更少空间。
若未安装 git，也可从出版社网站或以下链接手动下载源码压缩包并解压：
github.com/rasbt/reaso…

在终端进入 reasoning-from-scratch 目录。
在该目录中执行：

uv run jupyter lab

上述命令会启动 JupyterLab，你可以新建空白 Notebook 编写与执行代码，或直接打开第 2 章的 Notebook（其中包含本章覆盖的全部代码）。

TIP
可用 uv run script-name.py 运行 Python 脚本。
上述 uv run ... 会自动创建本地虚拟环境（通常在隐藏的 .venv/ 目录中），并从项目内的 pyproject.toml 自动安装依赖，因此无需再手动用 requirements 文件安装。如果你还需安装额外依赖，可执行：
uv add packagename
补充代码仓库在 ch02 子目录中还提供了更多安装说明与细节。

2.3 Understanding hardware needs and recommendations

（硬件需求与建议）

你可能听说过：训练 LLM 非常昂贵。对头部厂商而言，训练一个新的基础模型 LLM（尚未加入任何推理技术）在算力成本上往往少则百万美元量级、多则五千万美元以上。
这使得自行训练 LLM 对我和大多数读者而言都不现实。因此，本书将使用一个相对较小（但足够胜任）的预训练 LLM，并在其之上实现推理技术。
需要强调的是：这个更小的 LLM 在架构上与当代 SOTA 模型一致，只是缩小了规模；我们应用的推理方法与大型 LLM 上使用的方法相同。区别仅在于：小模型让我们可以以更可承受的成本来探索这些方法。

打个比方：如果你想学习汽车的工作原理，作为新手，你大概不会一上来就造一台昂贵的法拉利；更合理的是先从一辆“小甲壳虫”入手——同样能学到很多关于发动机和变速箱的知识，而且更容易把复杂的“细枝末节”先放到一边。
尽管本书为教学目的选用较小模型，但推理技术的使用、开发与应用仍然吃算力。在后续的第 5–7 章，如果能使用 GPU，会更有收益。

若你已按上一节安装了 PyTorch，可以用下面的代码检测你的机器是否有 PyTorch 支持的 GPU：

import torch

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

可能的输出示例：

PyTorch version 2.7.1
Only CPU

USING TENSOR CORES
（使用 Tensor Cores）

如果你有较新的 NVIDIA GPU（Volta 架构及以后），可以利用 Tensor Cores 来提升矩阵乘法吞吐。启用方式很简单：

torch.set_float32_matmul_precision("high")

默认情况下，PyTorch 运行在 "highest" 精度模式（不使用 Tensor Cores）。启用 "high" 可能会因浮点舍入而对结果产生微小影响，但就本书的实验来看，并不影响结果。
即便你没有 GPU，或不确定 GPU 是否支持 Tensor Cores，执行这行代码一般也是安全的。

不用担心没有 GPU：第 2–4 章 在 CPU 上也能在合理时间内运行。
根据不同章节，代码会自动优先使用可用的 NVIDIA GPU，否则回退到 CPU（某些部分或章节如推荐，将使用 Apple Silicon GPU）。相关信息会在相应章节进一步说明。

和许多每天与 LLM 打交道的研究者一样，我在家也没有足以训练 LLM 的 GPU 设备，通常使用云资源。若你在找云服务商，我个人偏好 Lightning AI Studio（lightning.ai/），因为其易用性与功能支持较好（见图2.4）。或者，Google Colab（colab.research.google.com/）也是不错的选择。

Figure 2.4 An overview of the Lightning AI GPU cloud platform in a web browser. The interface supports Python scripts, Jupyter notebooks, terminal access, and lets users switch between CPU and various GPU types based on their compute needs.
（Lightning AI GPU 云平台的浏览器界面概览：支持 Python 脚本、Jupyter 笔记本、终端访问，并可按需在 CPU 与多种 GPU 之间切换。）

就我写作时而言，Lightning AI 在注册与验证后还会提供免费算力额度，可用于图 2.4 展示的不同 GPU 选项。（本章不需要 GPU；如果你想用 GPU，L4 已经绰绰有余。）

NOTE
披露信息：我在 2023 年参与了 Lightning AI 平台的构建与发布，并仍持有少量权益。我未因推荐而获得赞助，平时也是自费使用。之所以使用它，是因为我觉得它最顺手：支持多种 GPU、可随时切换并回退到 CPU 以节省成本，还能暂停/恢复环境而无需重做设置。
补充代码仓库在 ch02 子目录中还提供了更多 GPU 平台推荐。

USING PYTORCH
（关于 PyTorch）

本节我们导入并使用了 PyTorch——当前最广泛使用的通用深度学习库。全书我们都将用它来运行与训练 LLM，包括后文将开发的推理方法。
如果你对 PyTorch 还不熟悉，建议阅读我的教程 “PyTorch in One Hour: From Tensors to Training Neural Networks on Multiple GPUs” （免费）：
sebastianraschka.com/teaching/py…

2.4 Preparing input texts for LLMs

（为 LLM 准备输入文本）

本节我们将演示如何使用 tokenizer 来处理 LLM 的输入与输出文本，如图 2.5 所示。图 2.5 在图 2.2 的“输入/输出准备”步骤基础上，进一步展开，为分词（tokenization）流水线提供更细致的视图。

Figure 2.5 A simplified illustration of how an LLM receives input data and generates output. The user-provided text is tokenized into IDs using the tokenizer's encode method, which are then processed by the LLM to generate output token IDs. These are decoded back into human-readable text using the tokenizer's decode method.
（LLM 接收输入并生成输出的简化示意：用户文本经 encode 转为 token IDs，LLM 处理后生成输出 token IDs，再由 decode 还原为可读文本。）

为了看看它在实践中的工作方式，我们先从本书的 reasoning-from-scratch Python 包中加载一个 tokenizer（你应已按第 2.2 节的说明完成安装）。
下载与本书所用 Qwen3 基础 LLM 对应的 tokenizer 文件，运行：

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)

该命令会下载约 6 MB 的 tokenizer-base.json，保存到 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)

由于我们尚未加载图 2.5 中的“核心组件”——LLM，我们先做一个只用 tokenizer 的干跑（dry run） ：进行一次分词往返（round-trip） ——先把文本编码为 token IDs，再把这些 IDs 解码回文本，如图 2.6 所示。

Figure 2.6 A demonstration of the round-trip tokenization process using a tokenizer. The user-provided input text is first converted into token IDs using the encode method, and then accurately reconstructed back into the original text using the decode method.
（使用 tokenizer 进行“往返分词”的演示：encode 将输入文本转为 token IDs，decode 将其准确还原为原始文本。）

图 2.6 底部所示编码过程的代码如下：

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

图 2.6 顶部所示解码过程（将 token IDs 还原回文本）的代码如下：

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

打印结果表明 tokenizer 能从 token IDs 重构原始输入：

'Explain large language models.'

在引入 LLM 之前，我们先看一下 encode 生成的 token IDs。下面的代码逐个打印 token ID 及其对应的解码结果，帮助理解 tokenizer 的工作方式：

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

输出如下：

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

如上所示，原始文本被切分为 6 个 token ID。每个 token 代表一个词或子词，取决于 tokenizer 的切分方式。
例如，“Explain” 被切成 “Ex” 与 “plain” 两个 token。这是因为该 tokenizer 采用了基于 BPE（Byte Pair Encoding） 的子词方法，既可用整词，也可用子词单元表示常见或罕见词。此外，token 中通常包含空格（如 " large"），帮助 LLM 感知词边界。

Qwen3Tokenizer 的词表规模约 151,000，在当前语境下算比较大（对比：早期 GPT-2 约 50,000，Llama 3 约 128,000）。
更大的词表会增加模型规模与单 token 的计算成本，但也能让更多词以单个 token 表示，而无需被拆分为子词。这通常是有利的：将一个词拆分（例如把 “Explain” 拆成 “Ex” 和 “plain”）会带来更多输入 token，从而拉长序列、增加耗时与资源。例如，token 数翻倍，模型的计算开销也大致会翻倍，因为模型需要生成更多 token 才能完成响应。

关于 tokenizer 的更详细讲解与从零实现，超出了本书的范围。若你感兴趣，可在附录 A 的延伸阅读中找到我的从零实现与更多资料。

EXERCISE 2.1: ENCODING UNKNOWN WORDS
试验 tokenizer 如何处理未知词：可以自创不存在的词；如果你懂多种语言，也可以尝试非英文词汇的编码。

2.5 Loading pre-trained models

（加载预训练模型）

上一节我们加载并熟悉了为 LLM 准备输入、并把 LLM 输出还原为可读文本的 tokenizer。本节将加载LLM 本体，如图 2.7 概览所示。

Figure 2.7 An overview of the four key stages in developing a reasoning model in this book. This section focuses on loading pre-trained LLM in Stage 1.
（本书推理模型开发的四个关键阶段概览。本节聚焦阶段 1：加载预训练 LLM。）

如前所述，本书使用 Qwen3 0.6B 作为预训练基础模型。本节我们将加载其预训练权重（见图 2.7）。“0.6B” 表示模型约含 6 亿个参数。

**为什么选择 Qwen3？**在评估数个开放权重（open-weight）的基础模型后，我选择 Qwen3 0.6B，主要基于以下考虑：

本书需要一个小而能打、且能在消费级硬件上运行的开源模型；
就我写作时点而言，Qwen3 系列的更大变体在模型效果上处于开源阵营的第一梯队；
相比 Llama 3 1B、OLMo 2 1B，Qwen3 0.6B 的显存占用更友好；
Qwen3 既提供基础模型（base） （本书推理方法的落脚点），也有官方推理变体，便于作为评测参照。

NOTE
“Qwen3”的规范写法中间不加空格；而“Llama 3”是有空格的。

秉持“from scratch”的精神，本书使用我用纯 PyTorch 重实现的 Qwen3，完全不依赖外部 LLM 库。此实现更强调可读性与可改造性，便于你做后续实验。尽管是从零实现，它仍然完全兼容原版 Qwen3 的预训练权重。
当然，本书不会深入讲解 Qwen3 的代码实现细节——这本身就足以写成另一本书（类似 Build A Large Language Model (From Scratch) ）。本书专注在基础模型之上的推理方法实现。

NOTE
该重实现的 Qwen3 完全本地运行，就像一般的 PyTorch 神经网络：无服务端组件、无外部 API 调用。所有模型使用都在你的机器上完成，数据不出本地。若你关注隐私，这样的设置能确保你对 LLM 输入与输出拥有完全控制。
关于 Qwen3 的更多实现细节与代码，可参考附录 C。

在加载模型之前，先指定将要使用的设备（CPU 或 GPU）。下面的函数会自动选择最佳设备：

def get_device():
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print("Using NVIDIA CUDA GPU")
    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("Intel GPU")
    else:
        device = torch.device("cpu")
        print("Using CPU")
    return device

device = get_device()

虽然 GPU 通常能显著提升速度与吞吐，但初次运行时，用 CPU 更有利于兼容性与调试。你也可以临时强制指定：

device = torch.device("cpu")

当你在 CPU 上验证本章代码无误后，注释或删除该强制设置并重跑；若系统具备 GPU，你应当会看到更好的性能。

NOTE
本章余下代码的基准是在 Apple M4 CPU 的 Mac mini 上执行的；章节末尾附有与 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)  #A
model.load_state_dict(torch.load(model_path))  #B
model.to(device)  #C
#A Instantiate a Qwen3 model with random weights as placeholders
#B Load the pre-trained weights into the model
#C Transfer the model to the designated device (e.g., "cuda")

注意：若 device 为 "cpu"，model.to(device) 实际不会执行设备迁移，因为模型默认就在 CPU 内存中。

运行上述代码后，你会看到类似下面的模型结构摘要（由 PyTorch 打印）：

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)
)

这段摘要展示了 Qwen3 0.6B 基础模型的核心组件：一个嵌入层、28 个堆叠的 TransformerBlock，以及最终的线性投影头。每个 block 含有 Grouped-Query Attention（GQA） 与前馈网络，并在多处使用 RMSNorm。
这些组件在图 2.8 中也有可视化示意。理解其细节并非本书必须——我们并不修改基础模型本身，而是在其之上构建推理方法。因此你现在可以把该结构当作一个“黑盒”。若感兴趣，可在附录 C查阅更多信息。

Figure 2.8 Overview of the Qwen3 0.6B model architecture. Input text is tokenized and passed through an embedding layer, followed by 28 repeated transformer blocks. Each block contains grouped-query attention, feedforward layers, and RMS normalization. The model ends with a final normalization and linear output layer. Arrows show the data flow through the model.
（Qwen3 0.6B 模型架构概览：输入文本经分词后进入嵌入层，随后通过 28 个重复的 transformer blocks；每个 block 含 GQA、前馈层与 RMSNorm；模型以最终归一化与线性输出层结束。箭头表示数据流。）

本节的关键信息：我们已经成功加载了一个可生成连贯文本的预训练模型（其架构如图 2.8）。下一节我们将编写文本生成函数，把已分词的数据送入模型，并以可读格式返回响应。

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

在加载了预训练 LLM 之后，我们的目标是编写一个利用 LLM 生成文本的函数。该函数将构成本书稍后实现的各类提升推理能力方法的基础，如图 2.9 所示。

Figure 2.9 本书中构建推理模型的四个关键阶段概览。本节解释 LLM 文本生成背后的主要概念，使我们能够在本章余下部分实现一个用于调用预训练 LLM 的文本生成函数。

在实现将贯穿本章与后续章节（见图 2.9）的文本生成函数之前，先回顾一下 LLM 文本生成的基本概念。
你也许已经知道，LLM 的文本生成是一个顺序过程：模型一次生成一个词（更准确地说是一个 token）。这通常被称为自回归（autoregressive）文本生成，如图 2.10 所示。

Figure 2.10 LLM 中顺序（自回归）文本生成的示意图。每一步，模型根据输入以及先前生成的 token 预测下一个 token；这些 token 会被累积回馈给模型，使其持续生成连贯的输出。

需要说明的是，图 2.10 仅给出了顺序生成过程的宏观概览。图中为了说明方便，展示的是给定一个输入提示词（prompt）后每一步只生成一个输出 token。
如果把其中任一次迭代放大来看，LLM 实际上会对每个输入 token都生成一个对应的输出 token。也就是说，如果我们有 6 个输入 token，模型会返回一个包含 6 个输出 token 的序列，如图 2.11 所示。不过，每次迭代我们真正关心的只有最后那个生成的 token。

Figure 2.11 放大观察自回归生成过程中的单次迭代。LLM 生成的输出序列与输入长度一致，但整体右移一位；在每次迭代中，模型预测序列中的下一个 token，从而实现对输入提示的逐 token 扩展。

在用图 2.11 的理念实现图 2.10 所示的整体自回归过程之前，我们先用第 2.4 节的示例提示词 “Explain large language models.” 再具体演示一次图 2.11 所示的细节：

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)  #A
input_tensor_fmt = input_tensor.unsqueeze(0)       #B
input_tensor_fmt = input_tensor_fmt.to(device)

output_tensor = model(input_tensor_fmt)            #C
output_tensor_fmt = output_tensor.squeeze(0)       #D
print(f"Formatted Output tensor shape: {output_tensor_fmt.shape}")
#A 将 Python 列表转为 PyTorch 张量
#B 增加一个维度（批维）
#C 前向计算得到输出
#D 去掉多余维度

SQUEEZING AND UNSQUEEZING TENSORS（张量的 squeeze 与 unsqueeze）

PyTorch 的 .squeeze() 与 .unsqueeze() 用于移除或添加长度为 1的维度，常见于为模型对齐输入形状。例如，模型可能期望二维输入（便于批量处理，见附录 E）。若当前只是 1D 行向量，可用 .unsqueeze(0) 在第 0 维增加一个维度：

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

输出：

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

这里 .unsqueeze(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 对应模型支持的词表大小。

直观理解：这 6 行中，每一行都是一个长度为 151,936 的向量，可视为对词表中每个 token 的“打分”；分值最高的位置就是本步要选取的生成 token（在这 151,936 个候选 token 中）。

因此，要得到“下一个生成词”，我们只需取这 6×151,936 矩阵的最后一行，在该行中找到最大分值对应的token ID，再用 tokenizer 将该 ID 转回文本，如图 2.12 所示。

Figure 2.12 放大观察单次生成迭代：如何把 LLM 输出的“原始分数矩阵”转换成一个 token ID 及其对应的文本形式。

下面用代码把图 2.12 的过程跑一遍。注意，LLM 是以“下一个词预测”为训练目标的；结合图 2.11，我们只关心最后一个 token，因此直接用 [-1] 取出：

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

此处 .detach() 用于切断该张量与梯度计算图的联系；简单说，我们只是拿到最后一行的分数向量用于解码，而无需保存反向传播相关的信息，既省内存又更高效。

打印结果是最后一个位置的 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 的 ID。再用 tokenizer 解码：

print(tokenizer.decode([20286]))

输出的生成 token 为：

 Large

MAX VS. ARGMAX（最大值与最大值索引）

PyTorch 中 torch.max() 返回最大值，torch.argmax() 返回最大值索引：

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

还可以用 keepdim=True 保持输出维度一致（便于与后续拼接等操作对齐）：

print(torch.argmax(example, keepdim=True))  # -> tensor([2])

小结一下：

图 2.10 说明了 LLM 的**迭代式（自回归）**文本生成；
图 2.11 放大观察了单次迭代；
图 2.12 进一步放大，展示了如何将分数矩阵转换为token ID与文本。

我们已经会用 LLM 生成单个 token。下一节我们将把这些概念串起来，实现一个顺序多步的文本生成函数，生成连贯的输出。

2.7 编写最小可用的文本生成函数

上一节解释了顺序文本生成中的单次迭代。本节在此基础上，编写一个文本生成函数：给定用户 prompt，驱动预训练 LLM 生成连贯的文本（参见本章概览中的图 2.13）。

Figure 2.13 本书构建推理模型的四个关键阶段概览。本节实现一个可用于预训练 LLM 的文本生成函数。

该函数（见图 2.13）的大致流程是：

将输入 prompt 转为 token IDs；
让模型预测下一个最可能的 token；
把该 token 追加到序列尾部；
用扩展后的序列再次喂给模型，预测下一个 token；
如此循环，直至满足停止条件；
最终把生成的 token IDs 解码为文本。

图 2.14 逐步展示了这一过程，并在每一步同时给出生成的 token ID 与其文本。（与前面的图 2.10 类似，但图 2.14 明确显示了 token ID。）

Figure 2.14 展示 LLM 的顺序（自回归）文本生成，并显式标注 token IDs。每次迭代，模型依据原始输入与此前已生成的全部 token 预测下一个 token，并将其以文本与ID两种形式加入序列。

下面的 Listing 2.1 实现了图 2.14 的顺序文本生成过程，使用上一节介绍的 argmax 选择策略：

Listing 2.1 一个基础的文本生成函数

@torch.inference_mode()                                        #A
def generate_text_basic(
    model,
    token_ids,
    max_new_tokens,
    eos_token_id=None
):
    input_length = token_ids.shape[1]
    model.eval()                                               #B

    for _ in range(max_new_tokens):
        out = model(token_ids)[:, -1]                          #C
        next_token = torch.argmax(out, dim=-1, keepdim=True)

        if (eos_token_id is not None                           #D
                and torch.all(next_token == eos_token_id)):
            break

        token_ids = torch.cat(                                 #E
            [token_ids, next_token], dim=1)                    #E
    return token_ids[:, input_length:]                         #F

#A 关闭梯度，节省显存/内存与加速推理
#B 切换到 eval 模式，保证确定性（最佳实践）
#C 仅取“最后一个位置”的分数用于解码
#D 若所有序列都生成了 EOS，则停止
#E 将新 token 追加到序列尾部
#F 仅返回“新生成”的 token（不含原始输入部分）

本质上，generate_text_basic 用一个 for 循环反复执行 argmax 选 token 的过程（由 max_new_tokens 控制迭代次数），返回如图 2.14 所示的生成 token IDs，之后再解码为文本。

我们用一个简单的 prompt 测试它是否工作正常（推理实验与更复杂的任务会在后续章节出现）。注意：下面的代码在许多机器上较慢，可能需要 1–3 分钟（我们会在后续小节显著加速它）：

prompt = "Explain large language models in a single sentence."
input_token_ids_tensor = torch.tensor(
    tokenizer.encode(prompt),
    device=device                            #A
    ).unsqueeze(0)

max_new_tokens = 100                        #B
output_token_ids_tensor = generate_text_basic(
    model=model,
    token_ids=input_token_ids_tensor,
    max_new_tokens=max_new_tokens,
)
output_text = tokenizer.decode(
    output_token_ids_tensor.squeeze(0).tolist()     #C
)
print(output_text)
#A 将输入 token IDs 放到与模型相同的设备（CPU/GPU）上
#B 令模型最多生成 100 个新 token
#C 将 PyTorch 张量转回 Python 列表以便解码

某次（CPU 上）生成的输出如下：

 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 用于训练中标记文档结束或区分样本。

TIP
输出首词 " Large" 前的空格源于：模型在延续输入继续生成，但我们在返回时用 token_ids[:, input_length:] 切掉了原始 prompt。如果你不想要这个空格，可以对返回文本调用 .lstrip()。

在推理时，我们通常希望模型在生成 <|endoftext|> 时立即停止。其 ID 为 151643，可通过下述代码确认：

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

该 ID 也保存在 tokenizer.eos_token_id 中。把它传给 generate_text_basic 即可启用遇 EOS 则停：

output_token_ids_tensor = generate_text_basic(
    model=model,
    token_ids=input_token_ids_tensor,
    max_new_tokens=max_new_tokens,
    eos_token_id=tokenizer.eos_token_id  #A
)
output_text = tokenizer.decode(
    output_token_ids_tensor.squeeze(0).tolist()
)
print(output_text)
#A 传入 EOS（序列结束）token 的 ID

输出类似：

 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.

与前一次相比，现在在遇到 EOS 后及时停止了。

你可能注意到：根据硬件不同（CPU/GPU），生成可能较慢。

EXERCISE 2.2: STREAMING TOKEN GENERATION（流式输出）
修改 generate_text_basic，在每次生成一个 token 时立刻返回/打印它，实现流式生成。

提示 1：用 yield 代替 return，将函数改为生成器。

提示 2：在函数外部对其进行迭代，逐 token 解码并打印（for token in generate_text_basic_stream(...): ...）。

在介绍加速方法之前，我们先写一个统计工具函数来测量生成时长：

def generate_stats(output_token_ids, tokenizer, start_time, end_time):
    total_time = end_time - start_time
    print(f"Time: {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():
            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()

    output_text = tokenizer.decode(output_token_ids.squeeze(0).tolist())
    print(f"\n{output_text}")

该函数会根据 start_time 与 end_time 计算总耗时、生成速度（tokens/sec）及GPU 峰值显存（目前仅支持 CUDA；PyTorch 尚未为 CPU 与 Apple Silicon 提供类似统计）。
使用方式如下：

import time

start_time = time.time()
output_token_ids_tensor = generate_text_basic(
    model=model,
    token_ids=input_token_ids_tensor,
    max_new_tokens=max_new_tokens,
    eos_token_id=tokenizer.eos_token_id
)
end_time = time.time()
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 tokens/sec 的速度看，确实比较慢。下一节我们将实现一种缓存技术，把生成速度提升5–6 倍。

TEXT GENERATION 与 INFERENCE 术语说明

在 LLM 文献或软件文档中，你常会看到 inference 一词，它常与“生成”互换使用。此处的 inference 源于神经网络语境：指用已训练好的模型进行预测（例如从一个 prompt 生成接下来的 tokens）。这与统计学中的 inference（对总体进行推断）不同。调用 generate_text_basic 时，我们就是在进行神经网络意义上的 inference。

2.8 通过 KV 缓存加速推理（Faster inference via KV caching）

现在我们已经有了一个基础的文本生成函数，可以看看它在实际运行时会发生什么。正如你可能注意到的，上节的文本生成有点慢。这也指向了一个关键关注点：推理阶段的性能。

当用 LLM 做推理（在本书语境下即从一个提示词生成文本）时，运行时性能（效率）会很快变得重要，尤其在长序列场景。虽然本书的代码更强调清晰性而非速度，但现实系统往往会采用一些工程技巧来让推理更高效。

在接下来的两节里，我们将介绍两种基本技术：KV 缓存与模型编译（见图 2.15 的总览），用以加速文本生成。

Figure 2.15 本书构建推理模型的四个关键阶段概览。本节在我们已加载的预训练 LLM 与前面编写的基础文本生成函数之上，引入 KV 缓存以提升执行速度。

如图 2.15 所示，提高文本生成速度的一种工程技巧是 KV caching。其中 KV 指注意力机制中使用的 keys 和 values。如果你对这些术语不熟悉也没关系。关键思想是：我们可以将某些中间结果缓存起来，并在每一步生成时重用（见图 2.16），从而加速推理。

Figure 2.16 KV 缓存在自回归文本生成中提升效率的示意图。与其在每一步都重新处理整段输入序列，不如将中间表示存入 KV 缓存，后续生成下一个 token 时直接复用，避免在每次迭代都把“新 token 与既有输入拼接后整段重算”。

图 2.16 所示 KV 缓存的核心思想，是在每次迭代时把计算得到的中间值存入缓存。此前，每生成一个新 token，我们都会把它与全量输入序列拼接后再次送入模型（图中的打叉框）。这种方式效率低，因为除了最新生成的 token 外，其余 token 在后续迭代中并未改变。借助 KV 缓存，我们可以避免冗余计算，直接取出已存的中间表示。

前文提到，诸如 KV 缓存这样的“非推理方法本体”的 LLM 细节，超出了本书重点，理解这些并不是后续章节所必需。不过如果你对此机制感兴趣，可以参阅我的免费文章《Understanding and Coding the KV Cache in LLMs from Scratch》（magazine.sebastianraschka.com/p/coding-th…）。

下面给出一个加入 KV 缓存的文本生成函数，它与清单 2.1 的基础版本几乎一致，不同之处已在注释中标出：

Listing 2.2 带 KV 缓存的基础文本生成函数

from reasoning_from_scratch.qwen3 import KVCache
 
@torch.inference_mode()
def generate_text_basic_cache(
    model,
    token_ids,
    max_new_tokens,
    eos_token_id=None
):
 
    input_length = token_ids.shape[1]
    model.eval()
    cache = KVCache(n_layers=model.cfg["n_layers"])         #A
    model.reset_kv_cache()
    out = model(token_ids, cache=cache)[:, -1]               #B
 
    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
 
        token_ids = torch.cat([token_ids, next_token], dim=1)
        out = model(next_token, cache=cache)[:, -1]          #C
 
    return token_ids[:, input_length:]
 
#A 初始化 KV 缓存
#B 首轮仍像之前一样将完整输入序列喂给模型
#C 之后的迭代仅将新生成的 next_token 作为输入（上下文从缓存取）

与清单 2.1 的 generate_text_basic 相比，generate_text_basic_cache 的主要区别是引入了 KVCache 对象。

在第一轮，我们仍把完整输入序列交给模型：model(token_ids, cache=cache)。在幕后，KV 缓存会为这些输入 token 逐层存储中间结果。
在后续迭代中，无需再传入整段序列，只需把刚生成的 next_token 喂给模型：model(next_token, cache=cache)。模型会从之前的 KV 缓存中取回上下文。

我们来计时看看加速效果：

start_time = time.time()
output_token_ids_tensor = generate_text_basic_cache(
    model=model,
    token_ids=input_token_ids_tensor,
    max_new_tokens=max_new_tokens,
    eos_token_id=tokenizer.eos_token_id,
)
end_time = time.time()
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.

可以看到，这种方式显著更快：从此前的 5 tokens/sec 提升到 29 tokens/sec（在 Mac Mini M4 CPU 上测得）。同时，生成文本与之前一致，这也是验证 KV 缓存实现是否正确的重要“健全性检查”。

下一节我们将介绍另一种进一步提升生成速度的技术，它会在后续评测章节派上用场。更快的生成意味着我们可以在更短时间内跑更多评测，更高效地对比不同模型或设置。

2.9 通过 PyTorch 模型编译进一步加速推理（Faster inference via PyTorch model compilation）

上一节我们介绍了 KV 缓存这一提速技巧，如图 2.17 的总览所示。

Figure 2.17 本书构建推理模型的四个关键阶段概览。本节在已实现的基础文本生成与 KV 缓存之上，进一步通过模型编译加速执行。

如图 2.17 所示，本节将使用 torch.compile 进一步显著加速推理。该特性允许对模型进行提前编译，以减少运行时开销并提升生成性能。

但在本书成稿时，torch.compile 在 MPS 设备（Apple Silicon GPU） 上支持不佳。对 Qwen3Model 使用它会引发 InductorError。
为保持跨设备兼容，我们先检测硬件类型，仅在支持的设备上启用编译：

if device.type == "mps":
    print("`torch.compile` is not supported"
          f" for the {model.__class__.__name__} model"
          " on MPS (Apple Silicon) as of this writing."
    )
    model_compiled = model  #A
else:
    major, minor = map(int, torch.__version__.split(".")[:2])
    if (major, minor) >= (2, 8):
        # 这可避免在 PyTorch 2.8+ 中，
        # 若模型包含类似 self.pos = self.pos + 1 的代码而反复触发重新编译
        torch._dynamo.config.allow_unspec_int_on_nn_module = True
 
    model_compiled = torch.compile(model)
#A 在 Apple Silicon GPU 上跳过编译

我们将编译后的模型赋给新变量 model_compiled，以保证本章其余代码无缝使用。
需要注意：首次执行编译后的模型可能更慢，因为要进行初始编译与优化。为更客观衡量性能提升，我们会重复执行多次文本生成。

首先，用未使用缓存的生成函数测试，代码与之前类似，但连续运行三次。执行可能需要几分钟：

for i in range(3):  #A
    start_time = time.time()
    output_token_ids_tensor = generate_text_basic(
        model=model_compiled,
        token_ids=input_token_ids_tensor,
        max_new_tokens=max_new_tokens,
        eos_token_id=tokenizer.eos_token_id
    )
    end_time = time.time()
 
    if i == 0:      #B
        print("Warm-up run")
    else:
        print(f"Timed run {i}:")
    generate_stats(output_token_ids_tensor, tokenizer, start_time, end_time)
    print(f"\n{30*'-'}\n")
#A 连续运行三次
#B 将首次运行标记为热身（编译/优化阶段）

输出示例：

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
Output text:
 
 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
Output text:
 
 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.
 
------------------------------

可以看到，编译后的模型有小幅提速：约 6 tokens/sec，相比此前的 5 tokens/sec。

接下来，用KV 缓存版本在相同流程下测试（仅把 generate_text_basic 换成 generate_text_basic_cache）：

for i in range(3):
    start_time = time.time()
    output_token_ids_tensor = generate_text_basic_cache(
        model=model_compiled,
        token_ids=input_token_ids_tensor,
        max_new_tokens=max_new_tokens,
        eos_token_id=tokenizer.eos_token_id
    )
    end_time = time.time()
 
    if i == 0:
        print("Warm-up run")
        generate_stats(
            output_token_ids_tensor, tokenizer, start_time, end_time
        )
    else:
        print(f"Timed run {i}:")
        generate_stats(output_token_ids, tokenizer, start_time, end_time)
 
    print(f"\n{30*'-'}\n")

输出示例：

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
Output text:
 
 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
Output text:
 
 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 上，带 KV 缓存的未编译模型为 29 tokens/sec，而编译后达到 68 tokens/sec，实现了超过 2 倍的加速。

EXERCISE 2.3：在非 CPU 设备上复现实验
如果你有 GPU，可在 GPU 上重跑本章代码，并将运行时间与 CPU 对比。

如果你关心 Apple Silicon GPU 与高端 NVIDIA GPU 下不同配置的对比，可参考表 2.1。

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

模式	硬件	Tokens/sec	GPU 显存
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	17	-
Regular compiled	Mac Mini M4 GPU	InductorError	-
KV cache	Mac Mini M4 GPU	18	-
KV cache compiled	Mac Mini M4 GPU	InductorError	-
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

如表中所示，NVIDIA GPU 的表现最好（符合预期）。不过在小模型与经优化的 CPU 端 KV 缓存实现下，CPU 的表现也相当不错。若模型更大或代码专门针对 GPU 优化，NVIDIA GPU 的优势会更明显。

同时要记住，输入序列越长，LLM 内部的注意力计算开销会按序列长度的平方增长，性能差异可能随之变化。

上述所有例子均在单条提示（批大小为 1）下运行。关于批量推理如何影响性能，参见附录 E。

2.10 小结（Summary）

使用 LLM 进行文本生成包含以下关键步骤：

环境搭建：准备运行 LLM 代码的开发环境并安装所需依赖。
加载预训练基础模型：加载预训练的基础 LLM（如 Qwen3 0.6B），后续章节将在其上扩展推理能力。
初始化并使用分词器（tokenizer） ：将文本输入转换为 token ID，并把模型输出解码回人类可读的文本。
自回归式文本生成：LLM 以**顺序（自回归）**方式逐 token 生成，通过预测“下一个最可能的 token”来扩展序列。
加速生成的两种手段：
- KV 缓存：缓存中间状态，避免在每一步对既有输入 token 反复重复计算。
- 模型编译（torch.compile） ：通过编译优化模型以提升运行时性能。

本章通过在预训练基础 LLM 上实现一个可用且高效的文本生成流水线，为后续章节中的推理能力构建打下了技术基础。