实现 GPT-2 的推理

本文目标说明

本文目标将用纯 TypeScript + node.js 实现对 gpt2-mini 推理，这是一个 FP32 精度、大小 147.26 MB，总参数量 38.60M 的微型模型，不依赖 python 生态、ONNX、Transfomer.js 等推理框架、也不涉及 GPU、SIMD 等通用硬件加速，整个推理管线完全架设在 v8 runtime 内来做，本文涵盖如下内容：

1. 模型文件格式及其解析加载
2. 张量与算子，张量编程初体验
3. GPT-2 的完整架构细节，包括 token 嵌入到最终输出 logits
4. 依据理论从零开始到实现完整推理，给出一个粗糙的模型推理引擎实现
5. 穿插介绍模型理论里出现的公式并解释它
6. 达成之前写的作品的闭环：

我将基于原生的 Float32Array 直接实现 GPT-2 内的所有推理计算，大概的几个清单数字如下：

实现前请降低预期，这个模型本身参数量极少只会输出一些阿巴阿巴的内容，而且 js 并不擅长高性能计算实际性能很差，每个 token 刚开始只跑几百毫秒到后面要跑数十秒，只需要明确目标：把理论和个人理解通过代码实现出来，这才是最重要的

"What I cannot create, I do not understand." —— Feynman

这里提供一个 check 清单，方便读者快速 review 一下自己是否适合阅读此文：如果对此清单内容有个 7 成以上的了解会比较适合本文，否则建议阅读下本人前作后再开始阅读本文：

首先第一步的分词直接用开源实现，因为这是纯算法内容不太涉及模型内部的计算，如下所示:

import gptTokenizer from 'gpt-tokenizer/cjs/encoding/r50k_base';
import { assert } from './assert';

/** 项目的 tokenizer，屏蔽下底层实现，只保留必要的接口 */
export const tokenizer = {
  encode: gptTokenizer.encode,
  decode: gptTokenizer.decode,
}

// "Hello": 15496, 见 gpt2 的 vocab.json
assert(tokenizer.encode('Hello')[0] === 15496, 'encode 编码错误');
assert(tokenizer.decode([15496]) === 'Hello', 'decode 解码错误');

// "hello": 31373, 见 gpt2 的 vocab.json
assert(tokenizer.encode('hello')[0] === 31373, 'encode 编码错误');
assert(tokenizer.decode([31373]) === 'hello', 'decode 解码错误');

模型文件：GGUF 与 safetensors

模型权重在磁盘上长什么样？社区目前最流行的两个主流格式是：safetensors 和 GGUF：

本文将选用 safetensors 的 gpt2-mini：FP32 精度、147.26 MB，约 38.60M 参数的微型模型，非常合适本场景

前向传播：推理就是一次张量流动

模型、或者说神经网络的推理，本质上就是个巨大的权重计算，从输入算到输出，称为前向传播，为了准确和高性能实现这个计算过程，需要引入张量的概念，用于建模和操作模型内的权重矩阵。

张量是什么 ?

张量 (Tensor) 听起来高大上，在数学里是更高阶的概念，用于统一标量、向量、矩阵等多维概念的建模，但在机器学习里更接近一种借用的命名，其本身在代码里更为朴素，用于在一段线性 Buffer 内建模和操作多维数据：

// ./src/tensor.ts
// 完整的张量比这个复杂的多
// 但这里我只需要一个最精简的核
export interface TensorF32 {
  readonly data: Float32Array;
  readonly shape: number[];
}

在张量视角下统一了多维结构的操作和抽象，因为模型文件内有大量的多维权重数据，实现基础的张量操作可以让我们直接在全局同一个 Float32Array 上直接跑计算过程而不是构造一大堆多维数组并且配套一大堆恼人的类型体操去处理各种边界细节。

张量与算子

在张量视角下的计算过程称为算子，比如下面这个简单的算子用来给整个张量都加上 10，如下所示：

// 通常算子的工程命名是 opXxxxx
function opAddScalar(x: TensorF32, c: number): TensorF32 {
  const data = new Float32Array(x.data.length);
  // 所有的张量底层都是一段线性内存，
  // 直接遍历做加法就能完成任意复杂的多维数据的全量加法
  for (let i = 0; i < x.data.length; i++) {
    data[i] = x.data[i] + c;
  }
  return { shape: x.shape, data };
}

模型的推理就是在写算子来实现论文里的公式，实现神经网络内的计算和输出结果，然而算子并不是一个直接对着公式照抄就能搞定的事，因为体系结构太复杂了：遍历一个数组从头到尾和从尾到头可能会有十倍的性能差距，写算子不仅是正确性的挑战，还得写出让 GPU/CPU 满意的代码，还需要直面大规模并发计算的挑战。

张量与神经网络的矩阵乘加

模型解决问题的方式是做输入输出的拟合，配合梯度逐渐调整拟合效果，最基础的线性拟合是一条直线 y = kx + b 但仅仅是一个输入是不够用的，需要引入更多的参数，比如三个参数 f(x_1, x_2, x_3)的线性映射可以这样展开：

\begin{aligned} Y = f(x_1, x_2, x_3) &= g ( w_1 x_1 + w_2 x_2 + w_3 x_3 + b_1 + b_2 + b_3 ) \\ & 因为 b_1 b_2 b_3 都是常数，可化简为 \\ \\ Y = f(x_1, x_2, x_3) &= g ( w_1 x_1 + w_2 x_2 + w_3 x_3 + b) \\ & 其中 g 是激活函数 \end{aligned}

上面可以画成下图的左部，而当我们进一步发散思维：如果输入 + 输出都是三个的时候则重复这个过程可以得到下图中的右部：

上图的右侧对应三个方程，共有 12 个参数，编号为 w{123|abc|ijk} 以及三个常数项 b{123} :

\begin{aligned} Y_1 &= g ( w_1 x_1 &+& w_2 x_2 &+& w_3 x_3 &+& b_1) \\ Y_2 &= g ( w_a x_1 &+& w_b x_2 &+& w_c x_3 &+& b_2) \\ Y_3 &= g ( w_i x_1 &+& w_j x_2 &+& w_k x_3 &+& b_3) \\ \end{aligned}

显然，这是一个带常数项的线性方程组，因此可以直接可写成两个矩阵来表达，W 代表齐次线性方程，B 代表一个常数项参数：

\begin{aligned} W &= \begin{bmatrix} w_1 & w_2 & w_3 \\ w_a & w_b & w_c \\ w_i & w_j & w_k \\ \end{bmatrix} & B = \begin{bmatrix} b_1 \\ b_2 \\ b_3 \\ \end{bmatrix} \\ \end{aligned}

那么，有关神经网络前向传播的定义，可以用一个矩阵运算直接给出，在形式上跟线性变换 y = kx + b 同构，极度优雅：

\begin{aligned} Y &= W X + B \\ \end{aligned}

将其展开就是这样的矩阵运算，其中的参数 W_\text{xxx} 最终就是「模型」的参数，利用这些参数来做预测，通常在训练开始的时候会设置为随机数，然后不断喂真实数据做训练找到损失函数的最小值：

\begin{aligned} \begin{bmatrix} Y_1 \\ Y_2 \\ Y_3 \\ \end{bmatrix} =& \begin{bmatrix} w_1 & w_2 & w_3 \\ w_a & w_b & w_c \\ w_i & w_j & w_k \\ \end{bmatrix} \cdot \begin{bmatrix} X_1 \\ X_2 \\ X_3 \\ \end{bmatrix} + \begin{bmatrix} b_1 \\ b_2 \\ b_3 \\ \end{bmatrix} \\ \end{aligned}

矩阵乘加是模型内运算量最大的运算，在张量视角下矩阵的乘法可以表达如下，通过张量表达运算过程中的参数、输入和输出：

GPT-2 模型架构和推理

GPT-2 的一次前向传播推理（forward）以及采样生成 nextToken 如下图所示：

一次推理后得到 nextToken 然后自回归重复调用生成：

读取 safetensors 加载张量权重

我们给出了张量的概念并介绍了其如何成为高性能推理最重要的心智模型，本章将会解析 gpt2-mini/model.safetensors 文件并将模型文件内存储的张量读进内存，转成一堆"可按名字取"的张量，本章还会深入展开张量的常用操作以及介绍 GPT-2 模型内的张量和概念，后文会用到。

张量的精简定义

PyTorch 的 torch.Tensor 等主流框架的张量实现有几十个字段：dtype、device、stride、grad、requires_grad、storage ... 但为了降低复杂度，只需要最精简的两个字段：

// src/tensor.ts
export interface TensorF32 {
  readonly data: Float32Array;
  readonly shape: number[];
}

data 是一段连续的内存空间， shape 是维度说明， 一个 [2, 3] 的矩阵和一个 [6] 的向量，底层可以是同一个 Float32Array，这种调整维数的操作称为 reshape。

两种张量：权重 vs 激活

推理时会存在两份不同的张量，一种是来自模型文件内的静态权重，另一种则是流动在模型网络中的激活张量，这些张量逐层传播流动应用算子并最终输出结果，用来保存临时结果，因此此处分别给出两个定义：

// 这份是激活张量，前面也有大量涉及了
export interface TensorF32 {
  readonly data: Float32Array;
  readonly shape: number[];
}

// 这份来自 model 模型文件直接解析出来
export interface ModelTensor {
  readonly dtype: 'F32';
  readonly shape: number[];
  readonly data_offsets: [number, number];
  readonly data: Float32Array;   // safetensors 文件的 subarray 视图
}

TensorF32 是 推理时算子之间流动的激活值，寿命短、每次算子产出一份新的； ModelTensor 是 只读权重， 多了一个 data_offsets，因为它的 data 并不是自己分配的， 而是从 safetensors 整块数据区里切出来的 subarray 零拷贝视图

操作张量的形状：reshape / split / transpose01

GPT-2 相对于现在的模型来说算比较精巧的，所以整套代码里和张量"形状"相关的操作只有三个， 它们构成了 Attention 里"多头拆分"那个看起来很唬人的操作的全部基石：

reshape 的零拷贝实现很好理解，但零拷贝的 split 和 transpose01 是否可行呢？可以，办法是给 TensorF32 引入新的 stride 字段用来描述每一维的"步长"，从而在逻辑上维持连续的张量视图，而底层物理字节未必连续，这是 PyTorch / NumPy 等框架的主流性能优化，可以省掉大量的张量拷贝操作，代价是所有算子取元素时都要写成 data[i*s0 + j*s1 + k*s2]，可读性下降一个量级，本文的目标是先跑起来再谈优化，所以先不做这个，只提供一个精简好懂的基础 TensorF32 即可。

note stride 本质上是"逻辑形状"和"物理字节"之间的一层间接索引，是推理性能优化的一个重要抽象； 它和 KV-Cache 共同面对的"内存墙"问题在后文会系统展开。

一个更容易理解的实例是如何描述一张 RGBA 图：最经典的建模是用一个 U32 表达一个像素，再通过位运算拆出 r g b a 四个分量，物理存储上是一段大的 Uint8ClampedArray（如 canvas.getImageData().data），字节序列是 RGBA RGBA RGBA … 这样交错排布，循环里靠 width × height 做二维定位、再用 +0/+1/+2/+3 取通道。把它套上张量视角，shape 就是 [H, W, 4]，类型 U8。

这个布局对 "按像素" 的算法天然友好，每个像素的 4 个分量挨着，但要做 "按通道并行" 的算法，比如直方图、亮度调整、灰度化、单通道卷积，其算法内的循环里要不停跨步，写起来啰嗦也不利于向量化，张量视角的解法是先做一次 transpose 换成 [4, H, W]，在这个视角下我们 “看到的” 是 4 行 H×W 块对应 4 个通道，遍历第一行的 H×W 就等于直接遍历了所有 r 通道了：在逻辑上我们实现了连续不跨步的遍历，而那些复杂的跨步边界细节由张量库在底层兜底了。

GPT-2 的张量清单

safetensors 的结构在已经图示过，落到代码上是四步，配合这四步可以读出来 GPT-2 的权重清单（张量清单）：

把上面这 4 步跑完后，可以得到一份张量清单，里面有详细的 GPT-2 模型权重细节：

header 解析出来就是一个 Record<string, 权重细节> 如下表格所示：

只有带权重的网络结构才会反映在上述清单内，模型内还有残差连接等其他无参数结构是不会反映在上述清单中的，后文会一一按顺序介绍，主要是：

读取模型的接口 class Model

提供class Model 用于直接读取 .safetensors 上的张量清单 (从 header 里读 & 解析)：

export class Model {
  public header: SafetensorsHeader;
  // 文件路径
  public constructor(filePath: string) {
    // parser 直接读取文件并做 JSON.parse 返回
    this.header = readSafetensorsHeader(filePath);
  }
  // name 形如 'transformer.ln_f.weight' 这样的字符串
  public getTensor(name: string): ModelTensor {
    const tensor = this.header.tensors[name];
    assert(!!tensor, `Tensor ${name} not found`);
    return tensor;
  }
}

GPT-2 整体架构

从模型文件内读取出张量，下一步就是写算子把 GPT-2 的架构实现出来，在具体实现前先看看 gpt2-mini 的几个超参数（可以理解为是配置文件），从 config.json 里能直接读出：

GPT-2 的网络可以拆成两层：外层是 forward 里串起来的主干管线，从 prompt 到 logits 输出，内部是 opTransformerBlock 处理 4 层 Transformer Block：

在前面我们已经通过开源的 gpt-tokenizer 实现了 prompt 转 tokens: number[]，现在正式进入模型内部处理词嵌入、位置编码以及归一化，即上图中的 opEmbed

词嵌入、位置编码、归一化

把输入 token 转词向量叫做 embedding，对应权重为 wte；为了让模型理解词句的前后顺序，需要对输入的序列做位置编码，叫做 position encoding，对应权重为 wpe。

词嵌入 + 位置编码

在正式开始，需要先谈谈 token 嵌入，也就是词嵌入，我们输入的单个 token 只是数字并没有多强的表达能力，必须将词句嵌入到高维向量空间内通过数据集训练学习得到语义更丰富的词向量，在 GPT-2 中，词向量的维度是 512，其直接存储在张量 wte 里，它的 shapes 是 [50257, 512]，其中 50257 是词表 vocab.json 的总数，即 wte 是一个每行 512 总行数 50257 的张量。

词嵌入是直接用 token 的数值直接在 wte 权重张量 上查表，取出词向量：

const model = new Model('./model.safetensors');
const wte = model.getTensor('transformer.wte.weight');
// "Hello" 对应的 token 是 15496
const vecHello: TensorF32 = {
  shape: [512],
  data: wte.data.subarray(15496 * 512, 15496 * 512 + 512),
}

位置编码呢？更暴力，直接用 token 在序列里的下标做另外一个权重的查表，叫做 wpe，如下图所示，如果输入的 prompt 只有 "Hello,World" 对应可以拆为 token 序列是 [ 15496, 11, 10603 ]，那么它的位置编码就是 [ 0, 1, 2 ]，然后将 0 1 2 在 wpe 张量上直接查表就能得到位置编码了。

总的来说，用 token 的数值做 wte 查表，用 token 的下标做 wpe 查表：

用公式表达 wte 和 wpe 非常简洁优雅，其中 i 是序列的下标编号，下列公式准确定义了用 token 的数值和下标分别对 wte wpe 进行查表操作，并将查表结果直接：

x_i \;=\; \mathrm{wte}\!\left[\, \mathrm{token\_ids}_i \,\right] \;+\; \mathrm{wpe}\!\left[\, i \,\right] \qquad i = 0, 1, \ldots, \text{len} - 1

完整的实现如下，当然也可以先做 wpe 拿到 [len, D] 位置编码张量，然后做 wte 查表直接加上去，这样可以减少访存：

// src/op_embed.ts
export function opEmbed(
  wte: ModelTensor,
  wpe: ModelTensor,
  tokens: number[],
): TensorF32 {
  assert(wte.shape.length === 2, `wte shape length must be 2`);
  assert(wpe.shape.length === 2, `wpe shape length must be 2`);
  assert(wte.shape[1] === wpe.shape[1], `wte wpe dim must be equal`);

  const [V, D] = wte.shape;
  const len = tokens.length;
  const data = new Float32Array(len * D);

  for (let i = 0; i < len; i++) {
    const token = tokens[i];
    assert(token < V, `Token ${token} out of range`);

    const wordBegin = token * D;
    const wordVec = wte.data.subarray(wordBegin, wordBegin + D);
    const positionVec = wpe.data.subarray(i * D, i * D + D);

    vecAddInto(wordVec, positionVec, data, i * D);
  }

  return {
    shape: [len, D],
    data,
  };
}

这里提一下词表有多大，gpt2-mini 总参数 38.6M，但 wte 一个张量就占了 25.7M，即 50257 * 512 = 25.7M，可以看到 67% 的参数都在词表嵌入里，这几乎是所有"小模型"的普遍遭遇：

另外 wpe 的形状中的 [512, 512] 第一维代表最大上下文（也就是说这份 wpe 只能给最大 512 的下标位置做位置编码，超过就是未定义行为了），所以在输出内容的时候基本就是阿巴阿巴输出点语法对但意思完全混乱的古神呓语 🐶，另外还需要注意小模型里 wte 占比太大了，而在更大规模的模型里，主要的参数量在集中在堆叠的 Transformer Block 中，如果一个走查表位置编码的模型能支持到 100K 上下文，那么对应的参数量将达到 100K * 512 = 51.2M，如果是 1M 上下文，则对应 512M 参数量，然而这对 1000B(1T) 总参数量来说算比较少了，因为最近新出的 DeepSeek、KIMI 都是 T 级大模（备注现在新的模型用的都不是按绝对位置做的位置编码了，通常是 RoPE / ALiBi 这类结构，所以实际的计算方式也会不同，此处仅做比方）

LayerNorm 归一化

模型内需要堆叠大量层数串在一起，每一层都会把前一层的输出（激活）相加、相乘、再送进下一层，每层持续激活得到流动的张量最终得到输出结果， 如果过程中对层的输入输出不做干预，数值幅度会随着层数推进而指数级漂移，上一层还在个位数量级，过两层就飙到百位数，再过两层就是数十万级了，LayerNorm 的作用就是在每层入口把数值漂移掐住：不管上游送进来什么，都先拉回一个标准的数值区间再往下走，并且还可以在这里引入额外参数做数值微调：

更进一步说，对于深层网络的数值稳定性来说归一化的重要性：

• IEEE 浮点精度会随数值变大而崩塌： 浮点数不得不品尝的特色就是数值巨大的时候，其精度会急剧下降，在大数值区间里相邻两个可表示值之间的间隔（ULP）会指数级变大， 数值一旦漂到 10^3 \sim 10^4 量级，FP8/FP16 甚至 FP32 都开始丢有效位了，而且这种误差会被后续每一层的激活传递持续累积放大，很快就没法信任输出结果了。    正如右侧代码所示 F16 已经没法区别 0.1 的精度了

  > arr = new Float16Array(2)
  Float16Array(2) [ 0, 0 ]
  > arr[0] = 1333.0
  1333
  > arr[1] = 1333.1
  1333.1
  > arr[0] === arr[1]
  true

• 前向激活稳定，反向梯度才能稳定，LayerNorm 直接作用在前向激活上，但通过"每层入口都把激活拉回 \mu=0,\; \sigma=1 的标准区间"，就间接把反向梯度的量纲也锁在了 O(1) 级别，训练过程因此能稳定地走过几十上百层的深网而不发散。

GPT-2 使用 LayerNorm 做归一化，并且把它放在残差相加之前，这种放法叫做Pre-LN；与之对应的是原版鼻祖 Transformer (Attention is all you need) 内采用的Post-LN（在残差相加之后），但 Post-LN 在深网里非常难训，所以目前主流的 decoder-only 模型（GPT-3 / Llama / Mistral）几乎一边倒地选了 Pre-LN。

LayerNorm 的数学定义如下，其中 μ 指的是均值，σ 指的是标准差，ε 是一个极小的常量，用于防止除零错误，\gamma 和 \beta 是可学习的参数，x 和 y 分别代表输入输出：

y = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \varepsilon}} + \beta

用换元法的视角可以轻松理解这个公式：标准化（standardization）+ 对角仿射变换（diagonal affine）：

\begin{aligned} 令 && \quad t &= \frac{x - \mu}{\sqrt{\sigma^2 + \varepsilon}} && 内层：标准化 \\ 可得 && \quad y &= \gamma \cdot t + \beta && 外层：仿射变换 \\ \end{aligned}

• 内层 (x - \mu) / \sqrt{\sigma^2 + \varepsilon}： 其中 \varepsilon 是一个极小常量（本模型里是 1e-5），工程上可以忽略它对结果的影响，它是为了防止 "这一行恰好所有元素相同" 的极端情况（此时 \sigma = 0）避免触发 divide-by-zero，相当于一道数值稳定的防线，且避免了 NaN 传染整个前向过程，把一组数据平移到均值 0、方差 1，在统计学里叫做 z-score 标准化。
• 外层 \gamma \cdot (\cdot) + \beta：这就是一个标准的线性变换 y = kx + b 的高维推广，不过这里的 \gamma 的值比较特殊，称为对角仿射变换
• 对角仿射变换是标准仿射变换 Y = W X + B 的特例，限定 W 为仅对角线有值的对角矩阵 \gamma， 此时 \gamma 只负责沿每个坐标轴独立缩放，偏置 \beta 负责整体平移：

  \gamma = \begin{bmatrix} \gamma_1 & 0 & 0 & \cdots & 0 \\ 0 & \gamma_2 & 0 & \cdots & 0 \\ 0 & 0 & \gamma_3 & \cdots & 0 \\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & 0 & \cdots & \gamma_D \\ \end{bmatrix} \qquad \beta = \begin{bmatrix} \beta_1 \\ \beta_2 \\ \beta_3 \\ \vdots \\ \beta_D \\ \end{bmatrix}

好在哪：只需要 2D 个参数量即可完整表达这个对角仿射变换，如果是完整的线性变换参数量则达到是 D²，在这个基础上赋予了 LayerNorm 依据梯度做自适应的能力，可以给输入的各个轴做放缩和偏移调整，你可以想象下面这个二维变换：

\begin{bmatrix} x_{变换后} \\ y_{变换后} \\ \end{bmatrix} = \begin{bmatrix} \gamma_1 & 0 \\ 0 & \gamma_2 \\ \end{bmatrix} \cdot \begin{bmatrix} x \\ y \\ \end{bmatrix} + \begin{bmatrix} \beta_1 \\ \beta_2 \\ \end{bmatrix}

这种模式最早由 Batch Normalization（Ioffe & Szegedy, 2015）提出，原论文给出的动机正是：强制把每层拉到 \mu=0,\; \sigma=1 会削弱表达力，所以再留一对可学习的 scale / shift，让网络在需要的时候能把标准化"撤回"回去（极端情况下 \gamma=\sigma,\; \beta=\mu 就完全恢复了原始分布），先把数值规整到统一量纲，再让网络自己调整。

对 [T, D] 输入的每一行独立跑一次上面这件事，就是完整的 opLayerNorm， 具体实现约 25 行，含三趟顺序扫描：算均值 → 算方差 → 归一化并做 \gamma / \beta 仿射，见源码 src/op_layer_norm.ts

Transformer Block：Attention / MLP / 残差

本章展开说明 Transformer Block 的结构，包括自注意力、MLP 和残差，其结构在前面有简单说明，这里引用下，即下图的右侧部分的 opTransformerBlock 是推理管线另外一个核心部分：

在 opTransformerBlock 里最主要的两个子结构是 opAttention 和 opMlp，如下图所示：

GPT-2 所涉及的算子

opTransformerBlock 内是模型内计算量最大的部分，这里完整罗列下 GPT-2 推理所需要的完整算子清单：

后面的行文我会尽量按上述算子展开，此外查表、reshape、split、transpose01这几个不是算子，而是对张量的形状/索引操作，会在上述算子内部被调用。以上共同构成了 gpt2-mini 推理需要的所有计算操作。

算子 opLinear：带 bias 的 y = xW + b

前面突然一下子灌输了一大波关于 GPT 的架构细节，相信读者已经懵了，这里按顺序开始一步步构建，这里先从调用频率最高的 opLinear 开始，它就是线性变换 y = xW + b，注意 GPT-2 这里顺序并不是 Wx，如下所示:

export function opLinear(
  x: TensorF32, // 输入
  weight: ModelTensor, // W
  bias: ModelTensor, // B
): TensorF32 {
  const [len, D] = x.shape;
  const [_D, H] = weight.shape;

  const data = new Float32Array(len * H);
  for (let i = 0; i < len; i++) {
    const rowBase = i * D;
    for (let j = 0; j < H; j++) {
      let sum = bias.data[j];
      for (let k = 0; k < D; k++) {
        sum += x.data[rowBase + k] * weight.data[k * H + j];
      }
      data[i * H + j] = sum;
    }
  }

  return { shape: [len, H], data };
}

tips 把 bias 作为累加器的初值是一个小妙招，相当于把 y += xW; y += b 两遍写成一遍，减少访存

为了直观理解上面三重循环到底在算什么，这里带入一组小整数走一遍 y = xW + 0（即 bias 全 0 的情形）。取 D = 1,\; H = 3,\; \text{len} = 4， 输入 x 是 4 个 token、每个 token 只有 1 维特征，权重 W 把这 1 维映射到 3 维：

// 把 x、W、bias 对应写成张量后
// 直接喂进去自己跑一次看看结果
opLinear(
 { shape:[4, 1], data: [2, 3, 5, 7] },
 { shape:[1, 3], data: [1, 10, 100] },
 // bias 全 0
 { shape:[3],    data: [0, 0, 0]    },
);

外层 i 扫 4 行、次层 j 扫 3 列、最内层 k 在 D = 1 上只累加一次，所以每个输出格子就是一次乘法： y_{ij} = x_{i,0} \cdot W_{0,j}。挨个写开：

\begin{aligned} i=0:\; & y_{0} = \begin{bmatrix} 2 \cdot 1 & 2 \cdot 10 & 2 \cdot 100 \end{bmatrix} = \begin{bmatrix} 2 & 20 & 200 \end{bmatrix} \\ i=1:\; & y_{1} = \begin{bmatrix} 3 \cdot 1 & 3 \cdot 10 & 3 \cdot 100 \end{bmatrix} = \begin{bmatrix} 3 & 30 & 300 \end{bmatrix} \\ i=2:\; & y_{2} = \begin{bmatrix} 5 \cdot 1 & 5 \cdot 10 & 5 \cdot 100 \end{bmatrix} = \begin{bmatrix} 5 & 50 & 500 \end{bmatrix} \\ i=3:\; & y_{3} = \begin{bmatrix} 7 \cdot 1 & 7 \cdot 10 & 7 \cdot 100 \end{bmatrix} = \begin{bmatrix} 7 & 70 & 700 \end{bmatrix} \end{aligned}

拼起来就是最终输出矩阵 y \in \mathbb{R}^{4 \times 3}：

// opLinear 的返回值
// data 是按行摊平的一维 Float32Array
{
  shape: [4, 3],
  data: [
    2,  20, 200,
    3,  30, 300,
    5,  50, 500,
    7,  70, 700,
  ],
}

最后一点是，上面代码是按三层循环变量的嵌套顺序写的，称为 ijk 序，通常性能不是最好的，它并不是 Cache 友好的，最佳实践是 kij 序，不过在 V8 Runtime 下，Float32Array 读取后还要装箱一次，收益没那么高，所以本文保留了最直观的 ijk 写法。

opMatMul：不带 bias 的版本

只做矩阵乘法与 opLinear 唯一的区别：内层累加器从 0 开始，即 y = xW + 0 前面那个带入数值的例子就是 bias 为 0 的例子。

matmulTransposeB：A · B^T

这是 opMatMul 的变种，在 GPT-2 里经常需要做类似 A · B^T 的计算，其中B^T 是转置后的 B 为了加速这个过程一次算出来而不需要先做转置在做乘法的两次算子，可以单独为这种模式专门写一个算子，一次到位搞定 A 乘以 B 的转置，代码里大部分细节跟前面乘法类似，注意张量转置的处理：

export function matmulTransposeB(
  A: TensorF32,
  B: TensorF32,
  scale: number = 1,
): TensorF32 {
  const [M, K] = A.shape;
  const [N, K2] = B.shape;
  const data = new Float32Array(M * N);
  for (let i = 0; i < M; i++) {
    const aRowBase = i * K;        // A 的第 i 行起点
    for (let j = 0; j < N; j++) {
      const bRowBase = j * K;       // B 的第 j 行起点（= B.T 的第 j 列）
      let sum = 0;
      for (let k = 0; k < K; k++) {
        sum += A.data[aRowBase + k] * B.data[bRowBase + k];
      }
      data[i * N + j] = sum * scale;
    }
  }
  return {
    shape: [M, N], data,
  };
}

单头自注意力

直接跳到"多头"之前，先把最基础的单头自注意力的数学定义讲一遍，输入是一批 token 的经过前面的嵌入和位置编码以及 LayerNorm 处理之后的词向量激活张量 X: TensorF32[T, d] 代表 T 个 token，每个 d 维，数学表达是 X \in \mathbb{R}^{T \times d}，核心公式就一条来自 Attention is All You Need：

\mathrm{Attention}(Q, K, V) \;=\; \mathrm{softmax}\!\left(\frac{Q K^T}{\sqrt{d_k}}\right) V

其中的三个矩阵由输入 X 各自过一次线性变换得到：

Q = X W_Q,\quad K = X W_K,\quad V = X W_V

把公式从里往外读，其实是四步：

1. 计算相似度 Q K^T：Q 的每一行是一个 token 发出的"我在找什么"的查询向量（Query），K 的每一行则是每个 token 暴露给别人看的"我是谁"的键向量（Key），两者做点积就得到一张 T \times T 的相似度矩阵 (QK^T)_{ij} 代表"位置 i 的 Query 与位置 j 的 Key 有多契合"，原理来自点积运算就是在算向量的相似度。
2. 缩放除以 \sqrt{d_k}： 点积的数值方差会随维度 d_k 线性增长，不除以 \sqrt{d_k} 的话 softmax 会被压成近似 one-hot，梯度消失
3. 归一化 \mathrm{softmax}： 对每一行做 softmax，把相似度转成"这 T 个位置上的概率分布"，加起来等于 1，每一行告诉我们："位置 i 做这次计算时，应该把多少比例的注意力分给其他各位置。"
4. 读出结果： V 的每一行是每个 token 如果你关注我，我给你什么内容 的值向量（Value），用上一步得到的概率作为加权系数，把所有 token 的 Value 混合起来，就是位置 i 最终的输出，一次按"相关度"做的加权平均。

总结：Attention = 一次可微的"查表 + 加权平均"，Q 发出问题、K 表达身份、V 提供内容，softmax 决定谁被听到多少，这套机制能工作的关键在于 Q/K/V 三个是依据投影矩阵 W_Q, W_K, W_V 算出来的，后者都是可学习的参数，通过训练而来。

而在 GPT-2 这种 decoder-only 模型里还要额外加一步：因果掩码（causal mask） 需要保证位置 i 只能看到位置 \le i 的 K / V—，在实现上是通过将上三角那一半的相似度置为 -\infty，配合 softmax 的数学性质保证这些位置的权重变成零

多头自注意力 opAttention

从单头走向多头：GPT-2 不是只做一次上面那套 Attention，而是同时做 8 次，把一条 512 维向量拆成 8 段各 64 维，让 8 组独立的 Q/K/V 分别去学不同的关系，直观来说就是词向量空间拆开来分别去学不同的关系，比如某一头盯指代消解、某一头盯句法依存、某一头盯长距离话题，最后把 8 个头的输出拼回 512 维，拼接很粗暴直接 concat ：

（a）QKV 一次投影就出来了

GPT-2 权重里把 QKV 合并在 c_attn.weight [512, 1536] 和 c_attn.bias [1536] 里了，从这两个 shapes 也可以看出 Q K V 各占 512，对应词向量维数，所以一次可以直接算和拆出：

const qkv = opLinear(x, cAttnW, cAttnB);     // [T, 3*512]
const [q, k, v] = split(qkv, 3);             // 各 [T, 512]

一次算出三次算出快，权重只读一遍，合并投影的操作叫做 fused QKV，GPT-2 的 c_attn 是它最早的代表

（b）多头：用 reshape 把 512 维拆成 8 份 64 维

多头注意力的直觉是：把 512 维劈成 8 个独立的 64 维子空间，让 8 组 Q/K/V 分别学习不同的关系，某一头负责指代消解，某一头负责句法依存，某一头负责长距离话题等等，当然这些概念完全是纯数学的高维向量空间内的魔法，不可能直接映射到人类的理解里，只是打比方，具体实现如下：

// 多头拆分: [T, 512] → [T, 8, 64] → [8, T, 64]
const qMulti = transpose01(reshape(q, [T, nHead, dHead]));
const kMulti = transpose01(reshape(k, [T, nHead, dHead]));
const vMulti = transpose01(reshape(v, [T, nHead, dHead]));

（c）单头内部：QKᵀ → mask → softmax → V

按代码写出来极度紧凑，GPT-2 分了 8 头注意力 (nHead=8)：

for (let h = 0; h < nHead; h++) {
  const qHead = getHead(qMulti, h);
  const kHead = getHead(kMulti, h);
  const vHead = getHead(vMulti, h);

  // scores = causalMask(Q @ K.T / √d)
  const scores = matmulTransposeB(qHead, kHead, scale); // [T, T]
  opApplyCausalMask(scores);

  // context = softmax(scores) @ V
  const probs = opSoftmax(scores);
  const context = opMatMul(probs, vHead);  // [T, dHead]

  // 把 context [T, dHead] 写回 merged 第 h 个 dHead 槽位
  for (let t = 0; t < T; t++) {
    const srcBase = t * dHead;
    const dstBase = t * dModel + h * dHead;
    merged.data.set(
      context.data.subarray(srcBase, srcBase + dHead),
      dstBase,
    );
  }
}

（d）合并多头：用 set + subarray 拼回 [T, 512]

前面配合张量的 reshape transpose01 等操作，已经把 QKV 分别拆成了 [T, 8, 64]，现在要拼回 [T, 512]，后再过一个 opLinear(merged, c_proj.W, c_proj.b)， 把 8 个头各自产出的信号"综合"成一条 512 维的输出，公式如下：

\mathrm{MultiHead}(X) \;=\; \mathrm{Concat}(\mathrm{head}_1, \ldots, \mathrm{head}_8) \, W^{O}

MLP：模型的知识层

MLP 指的是多层感知机，在定义上是多个全连接层串联且层间必须要有非线性激活函数做非线性拟合（这个世界上的非线性超乎你的想象）， 所以总的来说就是带激活函数的 opLinear，核心实现只需要 3 行，很简单吧，但它与注意力层存在某种同构：

export function opMlp(x, cFcW, cFcB, cProjW, cProjB) {
  const fc  = opLinear(x, cFcW, cFcB);      // [T, 512] → [T, 2048]  升维
  const act = opGelu(fc);                   // 逐元素 GELU
  return opLinear(act, cProjW, cProjB);     // [T, 2048] → [T, 512]  降维
}

Attention 与 MLP 的同构

这是 Anthropic 在 Transformer Circuits 里反复强调的观点：MLP 是一组 key-value 记忆（knowledge neurons），c_fc 的每一行可以当作一个 "key"，c_proj 的每一列则是对应的 "value"，c_fc(x) 给出的 2048 维激活就是 "x 对这 2048 个 key 分别有多像"；GELU 把负激活压到 0，c_proj 再把激活过的 key 对应的 value 加权读出。

GELU：一条 S 型曲线

GPT-2 用的 tanh 近似版，在 HuggingFace 里叫 gelu_new：

const SQRT_2_OVER_PI = Math.sqrt(2 / Math.PI);     // ≈ 0.7978845608
const COEFF = 0.044715;

export function opGelu(x: TensorF32): TensorF32 {
  const data = new Float32Array(x.data.length);
  for (let i = 0; i < x.data.length; i++) {
    const v = x.data[i];
    const inner = SQRT_2_OVER_PI * (v + COEFF * v * v * v);
    data[i] = 0.5 * v * (1 + Math.tanh(inner));
  }
  return { shape: x.shape, data };
}

关键坑：必须用这个近似版，PyTorch 的 torch.nn.functional.gelu 有两个实现 erf 版和 tanh 版，GPT-2 训练时用的是 tanh 版，数值在尾数上能对齐，如果用 erf 版每个位置会漂移 1e-5 左右，会出现误差逐层累积的情况。

残差流：把 Attention 与 MLP 组装成一层 Block

把 Attention 和 MLP 装到一起就是 opTransformerBlock，一眼看过去就是一个 Pre-LN 的 attention / MLP 双残差结构：

export function opTransformerBlock(x, model, layerIdx, nHead) {
  const prefix = `transformer.h.${layerIdx}`;

  // ── Attention 分支 ──
  const h1 = opLayerNorm(x, w(`${prefix}.ln_1.weight`), w(`${prefix}.ln_1.bias`));
  const attnOut = opAttention(h1, ...);
  const xMid = opAdd(x, attnOut);               // 残差 1：原始 x + 注意力增量

  // ── MLP 分支 ──
  const h2 = opLayerNorm(xMid, w(`${prefix}.ln_2.weight`), w(`${prefix}.ln_2.bias`));
  const mlpOut = opMlp(h2, ...);
  return opAdd(xMid, mlpOut);                   // 残差 2：xMid + MLP 增量
}

残差流的两个作用

残差的设计一眼看上去很粗暴，但它是现代神经网络最重要的一个里程碑，是深层网络架构的最佳实践：

• 训练阶段：梯度可以沿 "加号" 直接回传到底层，几百层的深网也不会梯度消失
• 推理阶段：哪怕这层学不到东西（attnOut ≈ 0）， 输出至少还等于输入，至少有个恒等映射作保底，模型可以"跳过"对输出无用的层

从残差流的视角来看，多层 Transformer Block 堆叠其实就是 "往残差流里注入一点增量"，逐步形成输出，残差流 = 初始嵌入（wte + wpe）+ 每一层贡献的增量之和。在后面 LM Head 的章节里，我们会把最终的残差流（再过一次 ln_f）和词表嵌入做点积，这等价于在问："算完所有层之后，每个位置的残差流，方向上最接近词表里哪个词？"

多层 Block 串起来

完成了 Transformer Block 后整个 forward 的构造会变得比较清晰了，就是一个多层循环，以及最终走一次 opLmHead 返回输出的 logits：

let x = opEmbed(wte, wpe, tokens);                        // 步骤 0
for (let i = 0; i < N_LAYER; i++) {
  x = opTransformerBlock(x, model, i, N_HEAD);            // 步骤 1~4
}
x = opLayerNorm(x, lnFW, lnFB);                           // 步骤 5
return opLmHead(x, wte);                                  // 步骤 6

最后的问题是 opLmHead，见下一章。

LM Head：用点积测方向

经过 4 层 Block 之后，得到模型在残差流里累积出来的 "对每个位置的最终预测"，我们把它记为 hidden，LM Head 要做的事只有一件：把这份 hidden 映射成对词表里 50257 个 token 的"原始分数"，做法也极其干脆，把 hidden 的方向和 wte 里每一个 token 的方向各做一次点积，也就是看 hidden 跟词表里每个 token 的方向有多一致，矩阵形式写出来就是一次和 wte 转置的相乘：

\mathrm{logits} \;=\; \mathrm{hidden} \cdot \mathrm{wte}^T

export function opLmHead(hidden, wte) {
  return matmulTransposeB(hidden, wte);     // [T, 50257]
}

weight tying: 为什么要复用 wte

注意这里的第二个参数是 wte，复用了输入端的词嵌入权重！GPT-2 没有独立的 lm_head.weight 张量，它和 wte 共享同一份 25M 参数。这个技巧叫 weight tying。

直观解释：既然输入端要把 "token id → 向量"，输出端要把 "向量 → token id 的分数"，分别对应 wte 的正反两个方向，用同一组张量既能嵌入又能测相似度，是非常优雅的。

logits 是 "方向的点积"

点积的代数意义是：

\mathbf{a} \cdot \mathbf{b} \;=\; \lVert \mathbf{a} \rVert \cdot \lVert \mathbf{b} \rVert \cdot \cos(\theta)

所以 logits[t][i] = hidden[t] · wte[i] 可以解读为："残差流在位置 t 给出的预测向量，和词表里第 i 个 token 的嵌入向量，方向有多一致"，通过这种方式关联词表并在训练时通过 cross-entropy loss 塑造 hidden[t] 的方向，让它指向下一个正确 token 对应的 wte[i]； 推理时反过来测量方向：训练是塑造方向，推理是测量方向：

注意如果完整计算 hidden 和 wte^T 的点积，实际上是在给 token 的每个位置做预测，实践中我们只需要取最后一行做预测就行了，不过我这里的实现是全算的，最终返回的 logits 张量 shapes 是 [T, 50257]

采样：softmax、温度、top-k、top-p

forward 返回 logits: [T, 50257]， 但我们下一步只需要最后一行，即只需要对当前序列的下一个 token 的预测，采样器的任务是：在这 50257 个分数里选一个 id 出来。

三种采样策略

export type SampleMethod =
  | { type: 'greedy' }
  | { type: 'top-k'; k: number; temperature: number }
  | { type: 'top-p'; p: number; temperature: number };

export namespace SampleMethod {
  export const Greedy = (): SampleMethod => ({ type: 'greedy' });
  export const TopK = (k = 10, temperature = 1): SampleMethod =>
    ({ type: 'top-k', k, temperature });
  export const TopP = (p = 0.9, temperature = 1): SampleMethod =>
    ({ type: 'top-p', p, temperature });
}

temperature 与 softmax

由于浮点数大数范围的精度和 Infinity 问题，工程实现 softmax 常需要引入一个减法来减缓 e^{大的} 的浮点数溢出：

\begin{aligned} softmax(z_i, T) &= \frac{e^{(z_i / T) - m}}{ \sum^K_{j=1} e^{(z_j / T) - m} } \\ 其中 \,\, m &= \max\limits_{1 \le j \le K} (z_j / T) \,\,\,\, 即 z_j 们经过 T 缩放后的最大值 \end{aligned}

这个跟原来的 softmax 是完全等价的，可以这样证明：

\begin{aligned} softmax(z_i, T) &= \frac{ e^{z_i / T}}{ \sum^K_{j=1} e^{z_j / T} } \\ &= \frac{ \phantom {xx} e^{-m} \cdot e^{z_i / T} \phantom {xxxx} }{ \phantom {xx} e^{-m} \cdot \sum^K_{j=1} e^{z_j / T} } \,\, (分子分母同时乘以 e^{-m} ) \\ &= \frac{e^{(z_i / T) - m}}{ \sum^K_{j=1} e^{(z_j / T) - m} } \end{aligned}

• T = 1：原始分布
• T < 1：更锐利、更确定，T → 0 时退化为 greedy
• T > 1：更平滑、更随机，T → ∞ 时退化为均匀分布

最后 top-p 是当前最常用的采样策略，它的设计哲学是：让候选集大小去适应模型的确信度， 而不是让用户去预估 k 应该设多大，当模型很确定时（"2 + 2 = "）候选词的数量非常少，很确定； 当模型很迷茫时（"她看着窗外，说"）候选词的数量自动变大。

自回归实现输出

所有算子、所有采样策略都准备好了，整个推理引擎的对外入口generate 不超过 50 行，核心循环只有 20 行：

export function generate(model, prompt, method, maxNewTokens) {
  const tokens = tokenizer.encode(prompt);
  const generated = [];

  for (let step = 0; step < maxNewTokens; step++) {
    if (tokens.length >= MAX_CONTEXT) break;        // 1. 上下文截断

    const logits = forward(model, tokens);           // 2. 完整前向
    const nextToken = sampleFromLogits(logits, method); // 3. 采样

    if (nextToken === EOS_TOKEN_ID) break;              // 4. 遇到 EOS

    tokens.push(nextToken);                             // 5. 追加到上下文
    generated.push(nextToken);
    process.stdout.write(tokenizer.decode([nextToken])); // 流式输出
  }

  return tokenizer.decode(generated);
}

注意第 5 步 tokens.push(nextToken)：生成的新 token 并追加到输入数组，下一轮 forward 的输入长度就变成了 T + 1，实现自回归输出：

朴素自回归 ≠ 最优实现：注意里每一步都要把 [T+1] 整段重算一遍，总成本是 O(T²)。现代推理框架用 KV-Cache 把这件事降到 O(T)，这是为什么 decode 速度能突破物理极限的核心技巧。 备注：我这里没有实现 KV-Cache，会在的总结里一起介绍。

最终运行效果

$ tsx ./src/forward.ts
prompt: "i want to sleep!!"
prompt tokens: [72, 765, 284, 3993, 3228]
sampling: {"type":"greedy"}

I’m not sure if I’m going to sleep at night, but I’m sure I’m going to sleep

总结：我没做什么

本章介绍两个我没做的重大优化（stride、KV-Cache）以及它们共同面对的内存墙、回顾推理的关键步骤和谈谈模型量化

我为什么没做 stride？

PyTorch / NumPy 的张量除了 shape 之外还带一个 stride 字段，用于说明"某一维走一步在底层数组里跳多少元素"。 配合它 transpose 转置操作可以给出 零拷贝 实现，只需要改 shape 和 stride，底层字节完全不动：

本文选择不实现 stride，一律 row-major contiguous，需要换顺序就真的拷贝一次。 原因有三：

• 降低实现复杂度：有了 stride，所有取元素的循环都要写成类似这样的读取 data[i*s0 + j*s1 + k*s2]，每个算子都会更复杂，减少不必要的噪声
• 关注张量流动本身：本文的对象是"推理流水线里张量的流动"，不是高性能张量库的实现，所以简洁的实现更重要
• 不过早优化：stride 是性能优化，其实不涉及逻辑机制，主要在多头注意力计算那块用到，可以将其中关键的转置 transpose 操作从一次 512*T*sizeof(F32) 字节的拷贝降到 0 字节拷贝实现

显然：推理性能的根本瓶颈是 "内存墙"

推理需要大量搬运和操作张量数据，其次是自回归生成里有用空间换时间的操作，比如 KV-Cache，会进一步加大访存压力，此外现代 GPU / CPU 里的算力（FLOPS）的增长速率远远超过 DRAM 带宽的增长速率，于是推理时往往不是算不完，而是搬不动：权重张量从显存 / 内存里流到计算单元的速度，跟不上计算单元的吞吐胃口。

内存墙早在 1995 年就被 Wulf 与 McKee 在 《Hitting the Memory Wall: Implications of the Obvious》 里精准预言了，他们观察到 CPU 速度每年以约 60% 的速度增长，而 DRAM 访问延迟每年只降低约 7%，这个 gap 一定会导致内存墙的出现，即 CPU 不得不需要等待内存数据， 30 年过去，现代 CPU 一次浮点运算一个周期约 1 ns、而一次 DRAM 访问约 100 ns，有两个数量级的差距并且还在继续扩大，所以内存墙的问题越来越严重。

业界常用一个叫 MBU（Model Bandwidth Utilization） 的指标衡量推理和训练时真实读取权重的带宽 ÷ 硬件理论峰值带宽，越接近 1 说明"带宽被用满了"，根据 Databricks 在 《LLM Inference Performance Engineering》 里公布的实测，Llama2 系列在 A100 / H100 上 decode 阶段的 MBU 通常在 50% – 70% 之间， 模型越大，越容易撞到 memory-bound，而且剩下那部分时间 GPU 的算力几乎是闲置的，OpenAI / Anthropic 虽然没有公开对外发布此类数据，不过 MLPerf Inference 的官方结果集也能侧面印证同一个量级。

总而言之，大模型的大张量推理和训练注定会撞到 mem-bandwidth 这个瓶颈，在 "内存墙" 的限制下实现最高的吞吐是所有训练/推理框架必须面对的核心命题。

量化的构造和解释

内存墙的问题直接的解法就是让每个权重更小，使其精度降低，这个过程叫做量化，然而关于精度的问题很大程度上需要借助类型论的视角，尤其是配合信息论的代数的视角才能理解量化的构造和解释，我们来做一个注意力训练：从代数视角看"类型"，注意到可以用枚举的方式定义类型：

// 共 256 个成员，需要 log2(256) = 8 bit，即 0x00 ~ 0xFF
enum Uint8 { 0, 1, 2, ... , 255 }

// 共 2 个成员，需要 log2(2) = 1 bit，即 0b0 / 0b1
enum Bool { true, false }

// 集中你的注意力：只有 1 个成员，需要 log2(1) = 0 bit（没错，这就是 null / void / unit）
enum Unit { onlyOne }

// 进一步集中：0 个成员，log2(0) 是为定义的，其渐近线是 -∞，因此代码里一旦"产出"这种值，程序就中断了（悖论：new 了一个对象剩余内存反而更多了？）
// 这也是 TypeScript 里 never 出现的语义：不可能发生的 dead code
enum Never { }

// 注意力涣散："一切合法取值" 需要多少 bit？ 对，unknown 本质上就是一段任意长度的线性内存空间
enum Unknown { ...所有可能 }

// 3 值呢? 显然 log2(3) ≈ 1.58 bit, 这就是所谓的 1.58 BitNet
enum _101 { -1, 0, 1 }

sizeof(type) = log2(成员数) 这个计算显然来源于信息论，继续展开：

FP32 / FP16 / INT8 / INT4 的区别，说到底就是这个类型的编码空间的容量，即允许的取值有多少个，从 232 个砍到 216、 再砍到 28、24，每砍一刀都是对信息的一次压缩，比如颜色 RGBA 可以表达为一个 U32 转为黑白能使其变成灰值 U8 在这个过程中会丢失大量信息，但人眼对灰度更敏感，所以多数情况下即便转为黑白人眼还是能准确辨认图片里的细节。

回到推理：一个权重张量里住着几亿、几十亿个 FP16 数字，但它们真的需要 65536 个不同的取值吗？实测结论是不需要，大部分权重分布非常集中，此时从 FP16 压到 INT8 / INT4 几乎不会伤到模型能力，在尽量不降智的前提下，可以给权重换一个"更短"的类型。

另外不仅是内存压力，更短的类型计算速度显然也会更快，让本地设备跑大模型成为了可能；量化在工程上演化出一整套的活跃研究领域，比如浮点数的各种数值细节、误差控制等等，展开够写好几篇博客了，这里就不细说了。

一个值得一提的极端案例是 BitNet b1.58， 它直接把权重量化到 {-1, 0, +1} 这个 3 值类型按前面"代数视角"，此时每个权重只需要log2(3) ≈ 1.58 bit，这就是它名字的由来，更讨巧的是：所有的乘法都退化成了"加/减/不动"，浮点运算都不用了，这使得它在访存和算力两条战线上都极具想象空间，是今年业界相当期待的一个前沿视角：甚至我已经在脑海里有了一个关于 mov + 位操作一次性在一个 U32 上直接处理 16 个权重的粗糙想法了 ...

回到本文：为了保持"最小实现"的初衷，示例代码仍然全程使用 FP32 做推理，但生产环境里的模型推理几乎都会使用量化，它是内存墙框架下性价比最高的一把刀。(蛐蛐一下：圈内特有的模型降智搞不好就是在偷偷灰度部分流量到量化模型 ...)

我为什么没做 KV-Cache

前面也提到 KV-Cache 了，这是推理框架提高吞吐性能的重要优化手段：

在 decoder-only + causal mask 的 GPT 架构下，attention 是唯一跨位置的算子， 即之前生成过的 token 对应的历史位置的 K/V 一旦算出永恒不变了，也就是可以缓存起来下次自回归重新用了。

朴素自回归每生成一个 token，配合自注意力机制都要把整段 token 再前向一遍，此时生成 N 个 token 的总成本是 O(N²)。 KV-Cache 把每层每个历史位置的 K、V 缓存起来，decode 时只算新 token 那一个位置的 Q， 再和累积的 K/V 做一次 attention，每步成本降到 O(N)，整次生成从 O(N²) → O(N)， 是数量级的改进。

我这里没有做 KV-Cache 的原因和 stride 一样：保证最小实现效果，这两个都是复杂度较大的优化措施，加进来会让核心的前向流水线被这类优化的工程细节淹没，不过既然点到了，这里给个 claude 帮写的伪代码来说明一下：

// 每层一份；跨 forward 调用持续累积
type LayerCache = { K: TensorF32 /* [T, D] */; V: TensorF32 /* [T, D] */ };
const caches: LayerCache[] = Array(N_LAYER).fill(null);

export function forwardWithKV(model, newTokens) {
  // ── 1. Embedding 只对新 token 做（而不是整段历史）──
  let x = opEmbed(
    wte, wpe, newTokens,
    /* posStart = */ caches[0]?.K.shape[0] ?? 0
  );

  for (let i = 0; i < N_LAYER; i++) {
    // ── 2. 在 Attention 里把历史 K/V 拼进来 ──
    const [qNew, kNew, vNew] = qkvProj(x, model, i);      // 只对新 token 算 Q/K/V

    const kAll = concat(caches[i]?.K, kNew);              // [T_prev + T_new, D]
    const vAll = concat(caches[i]?.V, vNew);
    caches[i] = { K: kAll, V: vAll };                     // ← 3. 写回缓存

    // 注意：Q 只有新 token 这一段，K/V 是累积的
    const attnOut = attentionCore(qNew, kAll, vAll);
    x = mlpBlock(opAdd(x, attnOut), model, i);
  }
  return opLmHead(opLayerNorm(x, ...), wte);              // [T_new, 50257]
}

三处核心改动：(1) 每层都挂一个持久的 { K, V } 缓存； (2) forward 只对新追加的 token 做 embedding 和 Q/K/V 投影； (3) attention 里把新的 K/V 拼到缓存后面，但 Q 只保留新段： 于是 QK^T 从 [T, T] 退化成 [T_{\text{new}}, T]，每步的 attention 成本从 O(T²) 降到 O(T)。

不过要注意：首次推理时 KV-Cache 还是空的，必须对整条输入 prompt 完整跑一次来构建 KV-Cache，这一步叫 prefill，对应的线上指标就是首个 token 返回耗时（TTFT, Time To First Token），之后生成新的 token 时就有缓存可以加速生成了，此时叫 decode，对应的指标是 每 token 耗时（TPOT, Time Per Output Token）。现代推理框架会把推理拆成前面提到的 prefill 和 decode 两个阶段，前者是典型的 compute-bound 一次 T×T 的大矩阵乘，吃算力，后者是典型的 memory-bound，每步只算一行和维护缓存，访存压力相比更大吃带宽。

完整的推理步骤

把前面内容串起来：

1. tokenize：tokenizer.encode(prompt)：把字符串转 token id 数组，我这里直接用了 gpt-tokenizer
2. embedding：opEmbed(wte, wpe, tokens)： 直接查 wte 表 + 和做张量加法的 wpe 位置嵌入，得到 x: [T, 512] 作为残差流起点，输入到后续层中。
3. 4 × Transformer Block：每层先 Pre-LN → Attention（c_attn 合并 QKV → 多头 → causal mask → softmax → @V → c_proj）→ 残差加 → Pre-LN → MLP（c_fc → GELU → c_proj）→ 残差加。
4. final LayerNorm：opLayerNorm(x, ln_f.weight, ln_f.bias) 把残差流最后一次拉回稳态。
5. LM Head：hidden \cdot wte^T 把每个位置的残差流和词表里每个 token 的嵌入做一次点积，得到 logits: [T, 50257]。
6. last logits row：只需要 logits[T-1]，最后一行代表着 nextToken 的预测，因此如果最后一步完整跑完整个上一步提到的点积运算，实际上就是给每个位置都做了预测。
7. sample：按 greedy / top-k / top-p 等采样策略挑选一个 token 出来，这里可以配置温度。
8. 自回归生成：回到第 2 步，直到遇到 EOS 或达到最大长度，关于 EOS 这类 Special Token 可参考 《Tool Use 具体是如何实现的？》

大模型推理目前都是这套骨架上，配合众多工程上的取舍和高性能优化去做的，目前仍然是个高速发展的领域。

结语

亲手写一遍和看文章的 "貌似读懂" 是两种完全不同的理解粒度：

• 对模型的实感：从"公式"到"手感"，从抽象概念到真实文件以及流动的张量激活，现在再看任何一个 Transformer 变种（GQA、MoE、MLA…），基本能从架构图读出细节架构的差异和能力的侧重
• 张量编程的思维范式：整个前向传播其实就是张量在算子间的流动，这是完全不同于命令式编程的心智模型，专门为大规模并行计算而生
• 采样和模型是两件事：模型输出的是 logits 而不是概率，关于前者的构造是精妙的：来源于对词表的点积计算，从而得到了 hidden 输出内容跟词表里的谁最像从而实现预测
• 最小 DEMO 的魅力：在这之上能精准的明确 KV-Cache、FP16 量化、GPU 硬件加速等 ai infra works 应该放置于何处以及明确其价值，最小 DEMO 的魅力在于骨架完整，所有想 "读过但没做过" 的优化都可以在这条基线上直接做 playground 开发。

下面是引用和受启发的资料文献：

1. Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762  — 自注意力机制原始论文
2. Radford, A. et al. (2019). Language Models are Unsupervised Multitask Learners (GPT-2). OpenAI · GPT-2 paper  — GPT-2 原始论文
3. HuggingFace. safetensors — Simple, safe way to store and distribute tensors. github.com/huggingface/safetensors  — 开源的实现，不过核心解析很简单 60 行能写完
4. Sennrich, R., Haddow, B., Birch, A. (2015). Neural Machine Translation of Rare Words with Subword Units (BPE). arXiv:1508.07909  — 关于 BPE 的我这里直接用了开源实现
5. Hendrycks, D., Gimpel, K. (2016). Gaussian Error Linear Units (GELU). arXiv:1606.08415
6. Ba, J. L., Kiros, J. R., Hinton, G. E. (2016). Layer Normalization. arXiv:1607.06450
7. Ioffe, S., Szegedy, C. (2015). Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. arXiv:1502.03167  — 可学习 scale/shift 的范式源头
8. He, K., Zhang, X., Ren, S., Sun, J. (2015). Deep Residual Learning for Image Recognition. arXiv:1512.03385  — 残差连接 / residual stream 的起点
9. Press, O., Wolf, L. (2016). Using the Output Embedding to Improve Language Models (Weight Tying). arXiv:1608.05859  — 关于 wte 复用的，没看完
10. Inan, H., Khosravi, K., Socher, R. (2016). Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling. arXiv:1611.01462
11. Geva, M. et al. (2021). Transformer Feed-Forward Layers Are Key-Value Memories. arXiv:2012.14913  — 关于 KV Memories 的，没看完
12. Elhage, N. et al. / Anthropic (2021–2024). Transformer Circuits Thread. transformer-circuits.pub  — Anthropic 的博文站
13. Wulf, W. A., McKee, S. A. (1995). Hitting the Memory Wall: Implications of the Obvious. ACM SIGARCH Computer Architecture News  — 内存墙原始论文

<|endoftext|>