Building an LLM Inference Engine from ScratchPart I / Module 1

Part I · Foundations

Module 1：从 n-gram 到 Transformer

把“语言建模”从统计计数一步步推进到 decoder-only transformer，并明确本教程为什么只做推理。

学习目标

用数学语言描述 next-token prediction，并区分 token、vocabulary、logits、probability。
理解 n-gram 模型的思想和局限：上下文太短、稀疏、无法泛化。
理解 attention 的核心直觉：当前 token 可以“看”过去哪些 token，并给它们不同权重。
认识 decoder-only transformer 的整体形状：embedding → 多层 block → LM head。
区分训练与推理，并解释本教程为何只实现 inference engine。

1.1语言建模：next-token prediction

语言模型的核心问题是：给定前面的 token，预测下一个 token。将输入序列记作 $x_1,x_2,\ldots,x_t$ ，模型给出如下条件分布：

P(x_{t+1}\mid x_1,\ldots,x_t)

在实际系统中，文本不会直接进入模型。Tokenizer 先将文本切分为 token，并将每个 token 映射为整数 id；模型的输出同样不是文本，而是一个长度等于词表大小 $|V|$ 的向量，称为 logits。对 Qwen2.5 而言 $|V|=151936$ ，这一维度将在 LM head 与 argmax kernel 中反复出现。

定义：logits

Logits 是 softmax 之前的未归一化分数。第 $i$ 个 logit 越大，模型越倾向于选择词表中的第 $i$ 个 token。概率由 softmax 得到：

p_i = \frac{e^{z_i}}{\sum_j e^{z_j}}

图 1-1：文本生成的基本数据流：文本 → token ids → logits → 下一个 token。

从 logits 到概率：softmax 的稳定写法

直接计算 $e^{z_i}$ 可能溢出，所以实际实现通常会先减去最大 logit：

p_i = \frac{\exp(z_i - m)}{\sum_j \exp(z_j - m)}, \qquad m = \max_j z_j

这个公式与原 softmax 等价，但数值更稳定。后面我们实现 attention 的 softmax 时会反复用到这个技巧。

1.2从 n-gram 出发：统计语言模型的极限

在深度学习流行之前，最经典的语言模型是 n-gram。它只看最近 $n-1$ 个 token，用计数估计下一个 token：

P(x_t \mid x_{t-n+1},\ldots,x_{t-1}) \approx \frac{C(x_{t-n+1},\ldots,x_t)} {C(x_{t-n+1},\ldots,x_{t-1})}

例如 bigram 只看前一个词，trigram 看前两个词。这种方法直观且易于解释，但局限明显：随着上下文增长，可能的 token 组合数量指数膨胀，绝大多数组合在训练语料中从未出现，其计数为 0，导致概率估计极度稀疏。

图 1-2：n 越大，可区分的上下文越多，但数据稀疏问题迅速变严重。

直觉

n-gram 的“记忆”是离散的：只有见过相同上下文才有可靠估计。神经网络的“记忆”是连续的：它把 token 映射到向量空间，让相似上下文可以共享统计强度。

1.3为什么需要神经序列模型？

神经语言模型的关键变化是：不再为每个字符串片段单独记一个计数，而是把 token 映射成向量，并用参数共享的网络计算上下文表示。同一个权重矩阵会被用于很多位置，这使模型可以从大量相似模式中泛化。

如果用矩阵表示，一个 token id $i$ 会先通过 embedding table 变成向量：

h_0 = E[i], \qquad E \in \mathbb{R}^{|V|\times d_{\text{model}}}

然后这个向量会经过很多层相同结构的 transformer block。每一层都更新 residual stream：

h_{\ell+1} = h_\ell + \operatorname{Block}_\ell(h_\ell)

这也是我们后续代码结构的主线：embedding、若干 decoder layers、final norm、LM head。

1.4Attention 的核心直觉

Attention 的想法可以用一句话解释：当前 token 不必只看固定长度的最近上下文，而是可以给历史 token 分配不同权重。 如果当前位置的 query 是 $q$ ，历史位置的 key 是 $k_j$ ，value 是 $v_j$ ，那么 attention 输出为：

\operatorname{Attention}(q,K,V) = \sum_j \alpha_j v_j,\qquad \alpha_j = \operatorname{softmax}\left(\frac{qk_j^\top}{\sqrt{d}}\right)

图 1-3：Attention 不是固定窗口，而是对历史 token 做加权读取。

在 decoder-only 模型中，attention 必须遵守因果约束：当前位置不能偷看未来 token。也就是：

\alpha_{t,j}=0 \quad \text{for } j>t

1.5Decoder-only transformer 的整体形状

GPT、Llama、Qwen 都属于 decoder-only transformer。它们的基本结构可以概括为：

图 1-4：Decoder-only transformer 的高层结构；后续模块会把每个方块拆成具体算子。

这里的每一个方块最终都会落到代码中：embedding 是一次查表，attention 是若干 matmul + softmax + 加权求和， MLP 是线性层与激活函数，LM head 是把 hidden vector 投影回词表大小的 logits。

1.6推理 vs 训练：为什么本教程只做推理？

训练和推理使用同一个模型结构，但工程目标完全不同。训练要计算 loss、反向传播、梯度、优化器状态；推理只做前向计算，并且要尽可能低延迟地生成 token。

维度	训练 Training	推理 Inference
目标	更新权重，使 loss 下降	固定权重，生成输出
计算	前向 + 反向 + 优化器	前向
内存	权重、梯度、优化器状态、激活	权重、KV cache、中间 buffer
性能核心	大 batch 吞吐	prefill 吞吐 + decode 延迟
本教程	不实现	完整实现

小结

本章将语言建模从基于计数的 n-gram 推进到 decoder-only transformer 的整体结构。核心结论有三点：模型的输出始终是一个 $|V|$ 维的 logits 向量，经 softmax 后才得到概率；n-gram 只能利用离散且出现过的上下文，而 attention 允许当前 token 按学习到的权重读取全部历史；本引擎仅执行前向计算，因此性能关注点是 prefill 吞吐与 decode 延迟，而非训练吞吐。下一章 Module 2 将把图 1-4 中的每个模块展开为具体算子，并给出对应的 shape。

思考与练习

基础实际实现 softmax 时为什么要先减去最大的 logit？不减会怎样？

当某个 logit 较大时， $e^{z_i}$ 可能溢出为 inf，进而使结果变为 NaN。减去 $m=\max_j z_j$ 后，最大的指数项为 $e^0=1$ ，其余均不超过 1，既避免溢出，又与原始 softmax 等价（分子分母同乘 $e^{-m}$ 相消）。attention 中的 softmax 采用同样的处理。

基础Qwen2.5 的词表大小 $|V|$ 是多少？它会在引擎的哪一步变成一个具体维度？

$|V|=151936$ ，即 LM head 输出 logits 的长度：hidden vector 经 LM head 投影后得到一个 151936 维向量。由于该维度较大，贪心解码时将完整 logits（f32 下约 608 KB/token）读回 CPU 的开销不可忽略，因此 xinfer 实现了 GPU argmax，仅将选中的 token id 读回。

进阶用一句话说清 n-gram 的“稀疏”问题，再解释神经网络靠什么缓解它。

n-gram 为每个具体上下文片段单独计数；上下文增长后，绝大多数片段在语料中从未出现，概率被估计为 0。神经网络不存储片段，而是将 token 映射到连续向量空间，相似上下文位置相近并共享权重，因此对未出现但相似的上下文也能给出合理预测，即实现泛化。

进阶decoder-only 模型里的“因果约束”是什么？不加会出什么问题？

因果约束指位置 $t$ 的 attention 只能访问 $j\le t$ 的 token，对未来位置令 $\alpha_{t,j}=0$ 。若缺少该约束，训练时模型可直接读取目标 token，学到的是复制而非预测；推理时按 token 自回归生成，未来 token 尚不存在。因此该约束保证了训练与推理的一致性以及生成的有效性。

挑战对照图 1-4，把 embedding、decoder block、LM head 分别对应到引擎里大致是什么计算。

embedding 是一次查表：以 token id 从 $E\in\mathbb{R}^{|V|\times d}$ 中取出对应行。decoder block 逐层更新 residual stream $h_{\ell+1}=h_\ell+\mathrm{Block}_\ell(h_\ell)$ ，每个 block 包含 RMSNorm、attention（matmul + softmax + 加权求和），再经一层 RMSNorm 与 SwiGLU MLP，这些后续均实现为独立的 HLSL kernel。LM head 是一次大矩阵乘，将最终 hidden vector 投影回 $|V|$ 维 logits，对应 xinfer 中采用 2D dispatch grid 的 linear kernel。

本教程选择推理，是因为它足够完整：你仍然需要理解 transformer 的所有核心算子、GPU 内存访问、同步、tokenizer、采样与部署；但又避免了训练系统中大量与优化器和分布式训练相关的复杂度。

Lab 1纸上追踪一个 tiny transformer

在真正写代码前，请先用纸和一个很小的 NumPy 程序追踪一个 tiny transformer。设词表大小为 5，hidden size 为 3，输入 token 为 $[1,2]$ 。你需要写出每一步的张量形状：

Embedding： $[2] \rightarrow [2,3]$
生成 $Q,K,V$ ：每个都是 $[2,3]$
计算 causal attention score： $[2,2]$
softmax 后乘 $V$ ：得到 $[2,3]$
LM head： $[2,3]\rightarrow[2,5]$
只取最后一行 logits，用 argmax 得到下一 token。

import numpy as np

V, H = 5, 3
tokens = np.array([1, 2])
E = np.random.randn(V, H) * 0.01

x = E[tokens]              # [seq=2, hidden=3]
Wq = np.random.randn(H, H)
q = x @ Wq                # [2, 3]

# 练习：补全 K、V、attention、LM head

思考与练习

基础解释 token、vocabulary、logits、probability 的区别。

token：文本被 tokenizer 切分后的最小单位，对应一个整数 id。
vocabulary：所有可能 token 的集合，大小记为 $|V|$ （Qwen2.5 约 15.2 万）。
logits：模型最后输出的未归一化分数向量，长度等于 $|V|$ ，每个分量对应一个候选 token。
probability：对 logits 做 softmax 后得到的概率分布，所有分量非负且求和为 1。

关系链：token id 进入模型 → 模型输出 logits → softmax → probability → 采样得到下一个 token id。

基础为什么 softmax 通常要先减去最大 logit？

因为 $e^{z_i}$ 在 $z_i$ 较大时会数值溢出。减去最大值 $m=\max_j z_j$ 后，所有指数的参数都 $\le 0$ ， $e^{z_i-m}\le 1$ ，不会溢出。数学上 $\frac{e^{z_i-m}}{\sum_j e^{z_j-m}}=\frac{e^{z_i}}{\sum_j e^{z_j}}$ 与原式完全相等，所以这是“免费”的数值稳定技巧。后续 attention 的 softmax 也用同样手法。

进阶举一个 n-gram 模型会失败、但 attention 有机会处理得更好的长距离依赖例子。

例如：“The keys that I left on the kitchen table this morning are …”。主语 keys 与谓语 are 之间隔了很长的修饰从句。bigram/trigram 只能看最近 1–2 个词，无法把 are 的单复数与远处的 keys 关联起来。attention 可以让 are 位置的 query 直接“关注”到很早的 keys，从而获得长距离主谓一致信息。n-gram 要覆盖这种依赖需要极大的 $n$ ，导致组合爆炸与数据稀疏。

进阶如果词表大小是 151936，LM head 每个 token 输出多少个 logit？这对 readback 有什么影响？

每个 token 输出 151936 个 logit。若用 f32，每步 logits 约 $151936\times4\approx608$ KB。若每生成一个 token 都把完整 logits 从 GPU 读回 CPU 再做 argmax，这 608 KB/token 的传输会成为 decode 的主要瓶颈之一。这正是后续 GPU argmax 优化的动机：greedy 时只在 GPU 上归约出最大 logit 的 index，只读回 4 字节。

挑战用 NumPy 完成 Lab 1，并打印每个张量的 shape；解释 causal mask 的作用。

关键 shape 链（词表 5、hidden 3、输入 2 个 token）：embedding [2,3] → Q/K/V 各 [2,3] → score [2,2] → softmax 后乘 V [2,3] → LM head [2,5]，取最后一行做 argmax。

causal mask 的作用：把 score 矩阵中“未来位置”（ $j>t$ ）置为 $-\infty$ ，softmax 后这些位置权重为 0，保证位置 $t$ 只能关注 $0..t$ 。这维持了自回归性质——生成第 $t+1$ 个 token 时不能偷看尚未生成的 token，否则训练与推理就不一致。