AI推理工程所需的数学（完整路线图）

在这篇文章中，我将为你拆解AI推理工程所需的核心数学知识。同时，我也会分享我亲身实践过的学习路线和资源。让我们直接切入正题。

1. 数字与精度

硬件如何存储数字，以及当它无法存储时会发生什么

起初我跳过了这部分，因为它看起来像是琐碎知识。后来我花了两天时间调试一个量化模型，它输出完全无意义的结果，而原因就藏在这里。尽早学习这部分内容。

1.1 浮点数的构成

一个浮点数由三部分组成：一个符号位、若干指数位和若干尾数位。指数位控制你能表示的数字大小范围，尾数位控制你保留的有效数字位数。

每种格式只是在这两项任务之间对可用位数进行不同分配，因此一旦你理解了这种权衡，就能预测任何格式的表现。

1.2 你会遇到的格式

1.3 舍入及其后果

每一次算术运算都会将结果舍入到最接近的可表示值，这引入了微小的误差。

在长序列的操作中，这些误差会累积，这就是为什么在FP16中求和长向量会偏离正确答案，以及为什么快速内核通常即使在输入精度较低时也会将累加总数保持在FP32中。

你还需要正确理解溢出和下溢。在FP16中，一个异常大的激活值可能会超过最大值而变成无穷大，从此之后所有下游计算都会被破坏。

1.4 机器精度与抵消

机器精度是接近1.0的两个可表示数之间的最小差异。它告诉你两个浮点数需要多接近，比较它们才不再有意义。

灾难性抵消发生在你减去两个几乎相等的数时。前导数字相互抵消，剩下的主要是舍入误差。

这在归一化层和方差计算中经常出现，教科书公式会减去两个大而相似的量，而数值稳定的版本则通过重新排列算术来避免这种情况。

1.5 稳定版 Softmax

Softmax 涉及对每个输入计算指数。如果某个输入接近 90，其指数会完全超出 FP32 的表示范围。

标准的解决方法是先减去所有输入中的最大值，再进行指数运算。由于 Softmax 对输入整体平移一个常数保持不变，这在数学上会得到相同的结果。

这种技术的通用形式称为对数-求和-指数（log-sum-exp），值得彻底掌握，而不仅仅是死记 Softmax 的修正技巧。它是稳定计算对数概率的基础，而 FlashAttention 内部使用的在线 Softmax 正是这一思想的变体。

2. 量化数学

用更少的比特存储权重

量化是内存和带宽节省的主要来源。其核心数学是一个线性映射，因此比论文标题所暗示的要更容易理解。

2.1 映射关系

你是在将一个浮点数值范围映射到一小组整数上。两个数字定义了这种映射。

缩放因子（scale）表示每个整数步长对应的实际值大小，而零点（zero point）表示哪个整数对应于实际的零值。要进行量化，你需要除以缩放因子，加上零点，然后取整。要进行反量化，则执行相反的操作。

几乎所有的量化方法都是选择这两个数字的不同策略。

2.2 对称与非对称

对称量化将零点固定为零，并围绕零点居中范围。其运算成本更低，且对于权重而言效果良好，因为权重往往已大致居中分布。

非对称量化允许零点移动，能更好地处理偏斜分布。经过ReLU后的激活值永不为负，因此对称范围会浪费一半的整数在从未出现的值上。

2.3 粒度

你可以对整个张量使用一个缩放因子，每个通道一个，或对固定大小的值组（如64或128个）各用一个。

更细的粒度能带来更好的精度，因为张量某一部分的少数极端值不再决定所有其他值的缩放范围。代价是需存储和处理更多缩放因子。INT4方法通常需要组级缩放才能保持足够的精度以实用。

2.4 选择范围

在量化前，你必须决定映射自的浮点范围。最简单的选择是观察到的最大值和最小值，但单个极端值会拉伸范围，使你关心的值几乎没有分辨率。

百分位裁剪刻意截断尾部。更谨慎的方法是搜索候选范围，并保留误差最低的那个，误差以平方误差或与原始分布的KL散度来衡量。

权重范围可以精确测量，因为权重是固定的。激活范围取决于输入，因此你需要通过模型运行一个小型校准集，并记录所见。

2.5 异常值

在Transformer中，少数激活通道携带的值比其他所有值大出许多倍。如果你理解了为什么一个大的值会破坏张量其余部分的有效分辨率，那么各种异常值处理方法的动机就变得显而易见了。

它们要么隔离异常通道并保持较高精度，要么在权重和激活之间转移难点，要么重新缩放通道以使范围更加均匀。

2.6 曲率方法简述

更强的权重量化方法利用输出对哪些权重最敏感的信息，这通过Hessian矩阵近似得到。模型几乎不依赖的权重可以粗略舍入，而模型高度依赖的权重则会得到补偿。

你不需要推导更新规则就能使用这些方法，但了解Hessian被用作敏感性估计，就能解释为何这些方法远超简单舍入。

3. 计算线性代数

相同的运算，以成本计而非解释

如果你已完成训练路线图，你就知道矩阵乘法的含义。对于推理，你需要知道它的成本。

3.1 计算矩阵乘法的成本

将一个 m 乘以 k 的矩阵与一个 k 乘以 n 的矩阵相乘，大约需要 2 乘以 m 乘以 n 乘以 k 次浮点运算，因为每个输出元素都是一个长度为 k 的点积，每一步包含一次乘法和一次加法。此外，你还必须从内存中读取两个输入矩阵，并将输出写回。

对于你模型中的各种形状，要习惯用手算来估算这两种成本。在剖析任何内容之前，能够说出“这一层大概需要这么长时间”，是整个领域中最为实用的一种习惯。

3.2 矩阵乘以矩阵与矩阵乘以向量的对比

矩阵乘以矩阵时，每加载一个字节的数据会执行大量算术运算，因为每个加载的值会在多个输出元素中被重复使用。而矩阵乘以向量时，每字节执行的算术运算很少，因为每个权重被加载后，仅使用一次便被丢弃。

为单个请求逐个生成 token 本质上属于第二种情况，这一事实在很大程度上解释了为何服务语言模型如此困难。

3.3 算术强度

算术强度是指操作次数除以移动的字节数。这是你首先需要检查的指标。

如果强度高，芯片的计算速率会成为瓶颈。如果强度低，内存带宽则成为限制因素，计算单元会闲置等待数据。大多数推理工作负载属于第二种情况，这也是为何大量优化工作聚焦于减少字节移动，而非减少算术运算。

3.4 张量在内存中的布局

你需要理解行主序和列主序布局、步长，以及张量连续性的含义。同时，还需明白为何转置矩阵在数学上无变化，却会消耗实际时间，以及为何形状会被填充至特定倍数，以便硬件矩阵单元能够使用。

3.5 分块

分块是指将大型矩阵乘法拆分为适合快速片上内存的小块，并在驱逐前尽可能多地重用每个块。涉及的数学计算是关于缓存容量及每块被重用次数的算术，这是你阅读的每个快速内核背后的核心思想。

3.6 低秩结构与稀疏性

矩阵的秩告诉你它实际包含多少独立信息。奇异值分解允许你用占用内存更少的低秩版本近似一个矩阵，这是低秩压缩和低秩适配器的基础。

适配器是对基础权重进行的一种低秩更新，因此理解秩能让你明白为何将其合并到基础模型中是精确的，以及为何同时服务多个适配器需要不同的方法。

对于稀疏性，实际需要了解的是，像每组四个中两个非零这样的结构化模式有硬件支持并能带来实际加速，而非结构化稀疏往往尽管减少了算术运算却运行得更慢，因为索引开销超过了节省的部分。

4. 性能建模

在测量之前估算运行时间

这一部分与机器学习的关联最少，却是我最常使用的部分。

4.1 屋顶线模型

如果将可达到的性能与算术强度作图，你会得到两条界限。

一条是代表内存带宽的斜线，当强度较低时它会限制你；另一条是代表峰值计算能力的水平线，当强度较高时它会限制你。你运行的每个内核都位于这两条上限之一的下面。

学会将工作负载置于该图表上，能立即告诉你应优化内存访问还是算术运算，从而避免你花费一周时间去优化一个从来就不是瓶颈的部分。

4.2 带宽受限估算

当受限于内存时，运行时间大致等于需要移动的字节数除以设备的内存带宽。

以下是一个此类估算的示例。一个70亿参数的模型以FP16格式存储大约占用14GB。为了生成一个输出令牌，你需要读取每个权重一次，因此将14GB除以GPU的内存带宽，就能得出在考虑其他因素之前，每个令牌所需时间的最低限。

如果用4位权重重新进行这个除法，这个最低限大约会降低四倍，这就是为什么即使算术运算不变，量化也能加速单请求生成的原因。

4.3 延迟与吞吐量

延迟是指单个请求完成所需的时间。吞吐量是指每秒完成的请求数量。增加批处理大小会提高吞吐量，但会恶化每个请求的延迟，因此你是在曲线上选择一个点，而不是寻找一个两者兼优的设置。

4.4 利特尔法则

利特尔法则指出，系统中的平均请求数等于到达率乘以每个请求在系统中平均停留的时间。

这是一个简洁的公式，它回答了一个非常实际的问题：要达到给定的吞吐量，需要同时有多少个请求在飞行中。一旦知道了这个数字，你就知道需要预留多少KV缓存内存。

4.5 排队行为

你不需要太多正式的排队论知识，但必须深刻理解一个结论。当利用率接近满负荷时，等待时间会无界增长，而非线性增长。

将服务器利用率从90%提升到95%，排队延迟的恶化程度远超这五个百分点所暗示的，这正是容量规划总是预留余地的原因。

4.6 阿姆达尔定律

通过优化流水线中某一阶段所能获得的加速比，受限于该阶段占总运行时间的比例。在着手处理某个仅占运行时间百分之三的任务前，先计算这一点是很有用的。

4.7 跨设备拆分

当模型被拆分到多个GPU上时，设备间需要交换中间结果。你需要大致了解诸如全归约（all-reduce）和全收集（all-gather）等集合操作的成本，以及这些成本如何随设备数量和消息大小增长。

值得理解的失败案例是，拆分模型后反而导致速度变慢——这发生在你为了减少并非瓶颈的计算而增加通信时。

5. Transformer 推理算术

将上述所有内容应用于实际工作负载

5.1 行为相反的两个阶段

处理输入提示是一次性完成的，每个令牌并行计算。每加载一个字节都有大量的算术运算，因此该阶段通常受计算限制，每个令牌的成本较低。

生成输出是一个令牌一个令牌地进行的，每一步都会读取整个权重集以产生一个令牌。每个字节的算术运算很少，因此该阶段受内存限制，每个令牌的成本要高得多。

几乎所有的服务优化都存在，因为这两个阶段对硬件的要求不同。

5.2 调整KV缓存大小

这个计算决定了你能同时服务多少用户，因此值得能够凭记忆完成。

总大小取决于层数、键和值头的数量、每个头的维度、每个存储元素的字节数、序列长度和批处理大小。将这些相乘，并包括一个因子2，因为你存储了键和值。

完成计算后需要注意两件事。大小随序列长度和批处理大小线性增长，并且当你量化权重时它不会缩小。对于长上下文服务，KV缓存比权重问题更大。

5.3 注意力开始占据主导的位置

注意力分数的计算量随序列长度的平方增长，而前馈层的计算量则随其线性增长。在序列长度较短时，二次项足够小，可以忽略不计，此时前馈层主导计算成本。

存在一个交叉长度，超过这个长度后，注意力成为主要开销。对于你所处理的模型形状，大致了解这个交叉点落在何处是很有价值的。

5.4 减少键和值头的数量

分组查询注意力和多查询注意力保留了全部查询头的数量，但在各组之间共享键和值头。如果你用减少后的键和值头数量重新进行5.2节中的KV缓存计算，内存节省量就会直接从算术中得出。

5.5 分页与碎片化

与其为每个请求预留一个连续的内存块以应对最大可能的序列长度，你可以随着序列的增长，以小的固定大小块来分配KV缓存。

相关的关键指标包括：每个序列最后一个部分填充块中浪费的空间量、预留内存中实际存放真实数据的比例，以及共享相同提示前缀的两个请求之间可以共享多少内存。

5.6 跨时间批处理

静态批处理收集一组请求，同时运行它们，并等待所有请求完成。当请求生成不同数量的令牌时，短请求提前完成，其槽位空闲，而最长请求继续运行。

连续批处理在旧请求完成时将新请求加入空闲槽位。关键在于槽位随时间保持的占用率，以及能恢复多少空闲容量。

6. 解码中的概率

采样设置对输出分布的影响

6.1 Softmax与温度

Softmax将原始分数转换为概率分布。温度在应用Softmax之前对分数进行除法。低于1的温度会增加最高选择与其他选择之间的差距，将概率集中在更少的令牌上。

高于1的温度则使分布变得平坦。当温度趋近于零时，会收敛到始终选择得分最高的令牌。

将这一过程视为重塑分布，而非一般的创造性设置，会更容易理解输出质量变化的原因。

6.2 截断尾部

Top-k 保留概率最高的 k 个词元。Top-p，亦称核采样，保留概率之和达到 p 的最小词元集合。Min-p 则保留概率高于最高词元概率一定比例的词元。

这三种方法都会丢弃分布的一部分，然后对剩余部分重新归一化，使其总和再次为 1。它们的区别在于如何决定截断位置，特别是截断是否根据模型在该步骤的置信度进行自适应调整。

6.3 熵与困惑度

熵衡量给定步骤中分布的分散程度，因此它可作为模型不确定性的合理信号，也是自适应解码策略的有用输入。

困惑度是数据集上平均负对数概率的指数形式。当检查量化或优化改动是否损害模型时，人们通常报告这一标准数值。

6.4 KL散度

KL散度衡量一个概率分布与另一个概率分布之间的距离。实践中，你用它来比较优化模型与原始模型在相同输入上的输出分布。较小的散度意味着你的加速是免费获得的。

较大的散度则表明，即使基准测试分数看似相近，你也已改变了模型的行为。

6.5 投机解码

一个小而快的模型会提前提出几个令牌，而大模型则一次性验证所有这些令牌。验证步骤采用拒绝采样，其设计确保最终接受的令牌严格遵循大模型自身本应产生的分布。这一保证是核心亮点，值得仔细阅读其推导过程。

预期的加速效果是计算得出的，而非猜测。它取决于草稿令牌被接受的概率、每步提出的令牌数量，以及两个模型之间的成本比率。如果不进行这一计算，很容易构建出比原始方案更慢的系统。

6.6 可复现性

浮点数加法不具备结合性，因此以不同顺序相加同一组数字会得到略有差异的结果。改变批处理大小会改变内核中数值求和的顺序，从而影响逻辑值的最后几位，偶尔也会改变被采样的令牌。

因此，相同的提示和相同的随机种子在不同批处理大小下可能产生不同的输出，这是由算术运算方式决定的，而非服务代码中的错误所致。

7. 基准测试的统计方法

避免自欺欺人的速度测量

大多数公开的推理基准测试，往往会在以下几个方面出现问题：

7.1 使用百分位数而非平均值

报告应包含中位数、第95百分位数和第99百分位数。平均值会掩盖那些表现不佳的请求，而用户实际感受到的是缓慢的尾部延迟，而非平均值。

请记住，估算第99百分位数需要大量样本。如果你只运行了五十次请求，你的p99仅为一两个数据点，不应作为测量结果引用。

7.2 分阶段报告

首字延迟衡量的是用户等待内容出现的时间，主要受提示处理和排队时间影响。每输出字延迟则衡量此后的流式传输速度，主要受内存带宽限制。同时，你还需要端到端延迟和整体吞吐量。

由于提示处理和生成阶段存在不同的瓶颈，单一的每秒令牌数指标可能掩盖几乎所有问题，因此这些指标应分开报告。

7.3 噪声

你需要掌握足够的基础统计知识，以便区分改进与运行间的随机波动。这意味着要对测量结果计算标准差，并采用某种形式的置信区间，这样当两次相同运行间的噪声达到5%时，2%的提升就不会被轻易宣称为胜利。

此外，应舍弃最初几次运行的数据。初始调用涉及编译、内存分配和缓存冷启动，若将这些数据混入平均值，会使整体表现看起来更差、噪声更大，而实际情况并非如此。

7.4 负载生成

无论先前请求是否已完成，都以固定速率发送请求，这能展示系统在无法跟上负载时的行为，包括队列如何堆积。

若保持固定数量的工作进程，每个进程在发送下一个请求前等待响应，则会掩盖上述问题，因为响应缓慢会自动降低发送速率。

需要特别警惕的具体陷阱是“协调遗漏”。如果负载生成器在等待缓慢响应时阻塞，那么在系统挣扎的这段时间内它并未发送请求，因此最差的延迟从未被记录下来。这会悄然使许多已发布的数据失效。

8. 离散数学与位级操作

阅读真实实现时会遇到的一些小细节

8.1 位操作

将两个4位值打包进单个字节，再通过掩码和移位将其取出。一旦写过一次便觉得简单，在此之前则令人困惑，所以务必亲手写一次。

8.2 2的幂与对齐

块大小设定为16、64或128，这与内存事务大小及矩阵单元形状有关。你需要掌握向上取整除法、填充至倍数，以及对齐要求的基本概念。算术本身微不足道，但其性能影响却不容小觑。

8.3 哈希与缓存命中率

检测两个请求是否共享提示前缀，通常涉及对令牌块进行哈希处理。你需要能够推理命中率，以及特定命中率对有效吞吐量的影响，同时还需了解用于存储共享前缀的树结构的基本原理。

8.4 复杂度表示法

足以察觉调度器或驱逐策略是否在执行随活动请求数量呈二次方增长的工作。你并非在撰写证明，而是在检查当并发请求达到一千时，系统是否依然能够正常运行。

我是如何学会这些的

我并非按上述顺序学会这些内容的，你大概也不会。以下是我建议现在开始学习的人所走的路径。

1. 先建立框架

Horace He 的文章《从第一性原理理解深度学习加速》将所有内容归类为计算受限、内存受限和开销受限。先读这篇，因为这种分类方法将是你组织其余知识的基础。

然后，打开计算器，逐步推导 kipply 的《Transformer 推理算术》。其中包含实际公式和实际数字，亲手做一遍算术，是“认得”与“真正掌握”之间的区别。

2. 了解硬件在做什么

Kirk 和 Hwu 合著的《大规模并行处理器编程》是标准的 GPU 编程书籍。即使你从不编写生产级内核，其中关于内存层次结构和分块的章节也会改变你对成本的思考方式。

GPU MODE 系列讲座（原称 CUDA MODE）免费开放，展示了人们如何推理真实内核，这是从书本中难以获得的。

至于 Hennessy 和 Patterson 的《计算机体系结构：量化研究方法》，不要试图通读。关于内存层次结构和性能测量的章节是屋顶线推理和阿姆达尔定律的来源，这些才是你需要的。

3. 浮点数

David Goldberg 的《每个计算机科学家都应了解的浮点运算知识》虽年代久远，却仍是该主题的最佳读物。内容密集，建议分两遍阅读，而非试图一次吸收。

然后，故意去“搞破坏”。编写短脚本，让 FP16 溢出，展示长求和中的误差累积，让朴素的 softmax 产生无穷大。十五分钟的实践比阅读教会我的更多。

4. 量化

阅读一篇关于高效推理量化的综述论文，以掌握相关术语，然后再去阅读你实际使用的方法所依据的论文。一旦掌握了缩放因子和零点，这些论文大多可以一口气读完。

我推荐的练习是，自己在一个权重矩阵上实现量化和反量化，测量重建误差，然后从每个张量一个缩放因子切换到每个通道一个缩放因子，观察误差下降。那个下午给我的收获比读那些论文还要多。

5. 服务系统

阅读FlashAttention论文时，重点在于分块和在线softmax的思路，而非CUDA的具体细节。它所做的，主要是减少内存流量。

关于PagedAttention的vLLM论文，最清晰地解释了为何KV缓存内存管理，而非原始计算能力，才是限制你能服务多少用户的关键因素。

投机解码的论文值得一读，尤其是其中的拒绝采样论证，那确实是概率论中一个巧妙的部分。

6. 测量

Brendan Gregg的《系统性能》与机器学习毫无关系，这正是它的价值所在。它教会你诚实地测量，认真对待尾部延迟，并避免常见的基准测试错误。

另外，查一下“协调遗漏”这个概念，读读你能找到的相关资料。它改变了我此后看待每一个延迟数据的方式。

让知识扎根的循环

阅读带我走了一段路，然后便不再奏效。真正起作用的是这样的：

第三步才是学习发生的地方。我的估算与测量之间的每一个差距，都源自某个我尚未了解的东西，而追根究底教会我的，比最初的计算本身更多。

可以略过的部分

超出基础范围的优化理论并不需要。收敛性保证、凸性和鞍点分析对训练很重要，而对于推理，模型的权重已经固定。

同样，你可以跳过大部分贝叶斯统计，以及除了利特尔法则和4.5节中的利用率警示之外的所有正式排队论。推导反向传播是好的背景知识，但在这里你不会用到。

那些看起来枯燥却不可省略的内容是浮点行为、算术强度和屋顶线推理、KV缓存大小，以及诚实的百分位数测量。这四样会不断出现。列表其余部分则是深度内容，当具体问题迫使你时再拾起也不迟。

最后一件我花了很久才接受的事。训练路线图主要是微积分和概率，而这份路线图主要是数字节和估算时间。

如果你原本期待第一种数学，却感觉自己误入了系统课程，那正是正确的感受，这个领域确实如此。