战胜 LLM 推理中的非确定性

可再现性是科学进步的基石。然而，从大型语言模型中获得可再现结果异常困难。

例如，你可能观察到多次向 ChatGPT 提问同一个问题会得到不同的结果。这本身并不令人惊讶，因为从语言模型获得结果涉及"采样"，这个过程将语言模型的输出转换为概率分布，然后概率性地选择一个 token。

更令人惊讶的是，即使我们将温度调至 0（这意味着 LLM 总是选择概率最高的 token，即所谓的贪心采样，因此理论上采样应该是确定的），LLM API 在实践中仍然不是确定的（参见过往讨论这里、这里或这里）。即使在自己的硬件上用诸如 vLLM 或 SGLang 这样的开源推理库运行推理，采样仍然不是确定的（参见这里或这里）。

那么，为什么 LLM 推理引擎不是确定的呢？一个常见的假设是浮点数非结合性和并发执行的某种组合导致了非确定性，取决于哪个并发核心最先完成。我们将这称为 LLM 推理非确定性的"并发+浮点数"假设。例如，最近的一篇 arXiv 预印本写道：

GPU 中的浮点数运算表现出非结合性，意味着 $(a + b) + c \neq a + (b + c)$，这是由于有限精度和舍入误差。这一特性直接影响 Transformer 架构中注意力分数和 logits 的计算，其中跨多个线程的并行操作可能根据执行顺序产生不同的结果。

你也可以在其他地方找到"并发+浮点数"假设被重复提及，比如这里（"有速度权衡，为了使端点快速，GPU 被使用，GPU 执行并行[非确定性]计算。任何现代 GPU 神经网络计算都会受到这种影响。"），或这里（"因为 GPU 具有高度并行性，加法或乘法的顺序在每次执行时可能不同，这可能会导致输出的细微差异"）。

虽然这个假设并非完全错误，但它没有揭示全貌。例如，即使在 GPU 上，重复运行相同数据的矩阵乘法也总是会得到按位相等的结果。我们确实在使用浮点数。我们的 GPU 确实有大量并发。那么为什么在这个测试中我们看不到非确定性呢？

A = torch.randn(2048, 2048, device='cuda', dtype=torch.bfloat16)
B = torch.randn(2048, 2048, device='cuda', dtype=torch.bfloat16)
ref = torch.mm(A, B)
for _ in range(1000):
    assert (torch.mm(A, B) - ref).abs().max().item() == 0

要理解 LLM 推理非确定性的真实原因，我们必须更深入地研究。

遗憾的是，即使定义 LLM 推理的确定性意味着什么也很困难。也许令人困惑的是，以下陈述同时都是真的：

GPU 上的某些核心是非确定的。

然而，语言模型前向传递中使用的所有核心都是确定的。

此外，LLM 推理服务器（如 vLLM）的前向传递也可以声称是确定的。

尽管如此，从使用推理服务器的任何人的角度来看，结果是非确定的。

在本文中，我们将解释为什么"并发+浮点数"假设没有抓住要点，揭露 LLM 推理非确定性背后的真正元凶，并解释如何战胜非确定性，在 LLM 推理中获得真正可再现的结果。

原罪：浮点数非结合性

在讨论非确定性之前，解释为什么会出现数值差异是有用的。毕竟，我们通常认为机器学习模型是遵循交换律或结合律等结构规则的数学函数。难道不应该有一个"数学上正确"的结果，我们的机器学习库应该提供给我们吗？

元凶是浮点数非结合性。也就是说，对于浮点数：

(0.1 + 1e20) - 1e20
>>> 0
0.1 + (1e20 - 1e20)
>>> 0.1

讽刺的是，打破结合律正是使浮点数有用的原因。

浮点数之所以有用，是因为它们允许"动态"精度级别。为了便于解释，我们将使用十进制（而不是二进制），其中浮点数的格式为 $\text{mantissa} * 10^\text{exponent}$。我们还将对尾数使用 3 位，对指数使用 1 位。

例如，对于值 3450，我们可以精确地将其表示为 $3.45 * 10^3$。我们也可以表示更小的值，如 0.486 为 $4.86 * 10^{-1}$。这样，浮点数使我们能够表示非常小和非常大的值。在科学中，我们可能会说浮点数允许我们保持恒定数量的"有效数字"。

如果你将两个具有相同指数的浮点数相加，看起来类似于整数加法。例如，123（$1.23 * 10^2$）+ 456（$4.56 * 10^2$）得到 579（$5.79 * 10^2$）。

但是当我们将两个具有不同指数的浮点数相加时会发生什么，例如 1230 和 23.4？在这种情况下，精确结果是 1253.4。但是，我们一次只能保持 3 位精度。浮点加法因此会丢弃最后 2 位数字，得到值 $1.25 * 10^3$（或 1250）。

此时，我们已经破坏了信息。注意，这可能在每次添加两个具有不同"规模"（即不同指数）的浮点数时发生。而添加具有不同指数的浮点数一直在发生。实际上，如果我们能保证永远不需要不同的指数，我们可以只使用整数！

换句话说，每次我们以不同的顺序添加浮点数时，我们都可能得到完全不同的结果。以一个极端例子，根据顺序，对这个数组求和有 102 种可能的不同结果。

import random

vals = [1e-10, 1e-5, 1e-2, 1]
vals = vals + [-v for v in vals]

results = []
random.seed(42)
for _ in range(10000):
    random.shuffle(vals)
    results.append(sum(vals))

results = sorted(set(results))
print(f"There are {len(results)} unique results: {results}")

# Output:
# There are 102 unique results: [-8.326672684688674e-17, -7.45931094670027e-17, ..., 8.326672684688674e-17]

虽然这是输出不完全相同的根本原因，但它并不直接回答非确定性来自何处。它没有帮助我们理解为什么浮点值以不同的顺序相加，何时发生，或如何避免。

答案在于核心如何实现。

为什么核心并不总是以相同的顺序添加数字？

如上所述，核心以不同顺序添加数字的一个常见解释是"并发+浮点数"假设。该假设认为，如果并发线程完成的顺序是非确定的，并且累积顺序取决于并发线程完成的顺序（例如使用原子加法），那么我们的累积顺序也将是非确定的。

令人困惑的是，虽然这可能导致非确定性核心，但并发（和原子加法）最终在 LLM 推理非确定性中完全无关！为了解释真正的元凶是什么，让我们首先理解为什么现代 GPU 核心很少需要原子加法。

何时需要原子加法？

通常，GPU 在许多"核心"（即 SM）上并发启动程序。由于核心之间没有内在同步，如果核心需要相互通信，这就会带来挑战。例如，如果所有核心都必须累积到同一个元素，你可以使用"原子加法"（有时称为"取并加"）。原子加法是"非确定的"——结果累积的顺序完全取决于哪个核心最先完成。

具体来说，假设你用 100 个核心缩减一个 100 元素的向量（例如 torch.sum()）。虽然你可以并行加载所有 100 个元素，但我们最终必须缩减到单个元素。一种方法是使用某种"原子加法"原语，其中硬件保证所有加法都会被处理，但不保证顺序。

这通常是人们所说的"非确定性"——你用完全相同的输入执行相同的核心两次，得到不同的结果。这称为运行间非确定性，其中你使用完全相同的依赖项运行相同的 Python 脚本两次，但得到不同的结果。

尽管并发原子加法确实会使内核具有非确定性，但绝大多数内核并不需要原子加法。事实上，在典型的 LLM 前向传播中，通常一个原子加法都不存在。

考虑到并行化归约可以从原子加法中获益，这一点可能令人意外。原子加法最终并非必需，主要有两个原因。

沿“批次”维度通常已经存在足够的并行度，因此我们无须沿归约维度进行并行化。例如，假设我们不是对单个 100 维向量进行归约，而是并行归约 500 个向量。在这种情况下，我们可以在每个核心中归约完整的向量，并让每个核心处理不同的向量。

随着时间推移，大多数神经网络库都采用了各种策略，在不牺牲性能的情况下实现确定性。例如，我们可以执行“分割”（或树形）归约，将包含 100 个元素的归约拆分为五个包含 20 个元素的归约，从而实现五路并行。然后，为了合并剩余的五个元素，我们可以执行一个单独的“收尾”归约（它不进行并行化，但由于处理的元素足够少，成本很低），也可以使用信号量（确保每个并发线程块都以确定的顺序进行累加）。信号量策略在这里有相关描述。

由于这两个因素，对于绝大多数神经网络操作而言，避免使用原子加法所带来的性能损失微不足道。

仍有少数常见操作在避免使用原子操作时会产生显著的性能损失。例如 PyTorch 中的 scatter_add（a[b] += c）。不过，在 LLM 中常用的此类操作只有 FlashAttention 的反向传播。有趣的是：你知道广泛使用的 Triton 版 FlashAttention 反向传播实现在算法上其实与 Tri Dao 的 FlashAttention-2 论文不同吗？标准 Triton 实现会在反向传播中执行额外的重新计算，从而避免原子操作，但代价是 FLOPs 增加 40%！

然而，LLM 的前向传播不包含任何需要原子加法的操作。因此，LLM 的前向传播实际上具有“运行间确定性”。

Wikipedia 写道：“确定性算法是指这样一种算法：给定特定输入，它总会产生相同的输出。”在这里，只要给出完全相同的输入（即推理服务器正在处理的请求完全相同），前向传播就总会产生完全相同的输出。

但是，前向传播本身具有“确定性”，并不足以保证包含它的整个系统也具有确定性。例如，如果我们请求的输出依赖于并行处理的其他用户请求（如 batch-norm），会怎么样？由于每个请求都无从得知与其并行的请求是什么，因此从该请求的视角来看，整个 LLM 推理同样具有非确定性！

事实证明，我们请求的输出确实依赖于并行处理的其他用户请求。这并不是因为我们以某种方式在批次之间泄露了信息，而是因为前向传播缺乏“批次不变性”，导致请求的输出取决于前向传播的批次大小。

批次不变性与“确定性”

为了说明批次不变性，让我们简化系统，只关注矩阵乘法。你可以假设所有矩阵乘法实现都具有“运行间确定性”。严格来说，这并不完全正确，但大多数常见的矩阵乘法实现确实具有这一性质。然而，它们不具备“批次不变性”。换言之，当批次大小发生变化时，批次中的每个元素可能会得到不同的结果。

从数学角度来看，这是一种相当反常的性质。矩阵乘法对于批次中的每个元素都应该是“独立”的——批次中的其他元素以及批次大小，都不应影响批次中特定元素的计算结果。

然而，正如我们可以通过实验观察到的那样，事实并非如此。

import torch
torch.set_default_device('cuda')

B = 2048
D = 4096
a = torch.linspace(-1000, 1000, B*D).reshape(B, D)
b = torch.linspace(-1000, 1000, D*D).reshape(D, D)
# Doing a matrix vector multiplication by taking
# the first element of the batch
out1 = torch.mm(a[:1], b)
# Doing a matrix matrix multiplication and then taking
# the first element of the batch
out2 = torch.mm(a, b)[:1]
print((out1 - out2).abs().max()) # tensor(1669.2500, device='cuda:0')

请注意，这具有“运行间确定性”。如果多次运行该脚本，它每次都会以确定性的方式返回相同结果。它不具备“硬件/软件版本不变性”——你的 GPU 或 PyTorch 版本可能会返回不同的值，但在同一环境中，它应该始终以确定性的方式返回相同的值。

然而，当不具备批次不变性的内核被用于更大的推理系统时，该系统就可能变得具有非确定性。当你向推理端点发起查询时，从用户的视角来看，服务器当前承受的负载实际上是“非确定的”。负载决定了内核运行时的批次大小，因而会改变每个请求的最终结果！

如果将内核对某个属性不具备不变性（即批次大小），与该属性自身的非确定性（即服务器承受的负载）组合起来，就会得到一个非确定性系统。

换言之，几乎所有 LLM 推理端点具有非确定性的主要原因，是负载（以及由此产生的批次大小）会以非确定的方式变化！这种非确定性并非 GPU 所独有——由 CPU 或 TPU 提供服务的 LLM 推理端点同样存在这一非确定性来源。

因此，如果想在推理服务器中避免非确定性，就必须让内核实现批次不变性。为了理解如何实现这一点，我们先来看一下内核最初为何不具备批次不变性。

如何让内核具备批次不变性？

为了让 Transformer 实现具备批次不变性，我们必须让每个内核都具备批次不变性。幸运的是，我们可以假设每个逐点操作都具备批次不变性。尽管对于 PyTorch 等框架中的所有内核而言，这一点成立，但它并非天然成立。例如，CPU 上的一些内核实现会对数组的某些部分使用向量化内部函数，而对其他部分使用非向量化内部函数，这些内部函数的数值结果未必总能做到逐位完全相同。因此，我们只需关注涉及归约的三个操作：RMSNorm、矩阵乘法和注意力。与并行相关的归约不在本文的讨论范围内，但同样的原则依然适用。一个可能有用的小知识是：在 Blackwell 上，以及使用 CUDA 12.8 及更高版本的 Hopper 上，NVLink-Sharp 的交换机内归约是确定性的。与许多事情一样，这一信息可以在 NCCL 的 GitHub issue 中找到。

方便的是，这些操作也恰好按难度从低到高排列。为了在实现批次不变性的同时保持合理的性能，每个操作都需要额外考虑一些因素。我们先来讨论 RMSNorm。

具备批次不变性的 RMSNorm

# x: [batch_size, hidden_dim]
# weight: [hidden_dim]
def rms_norm(x, weight):
    return x * torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True)) * weight

批次不变性要求，无论内核的批次大小如何，每个元素的归约顺序都必须固定。请注意，这并不意味着我们必须始终使用同一种归约策略。例如，如果被归约的元素数量发生变化，即使归约策略也随之变化，我们仍然可以保持批次不变性。Quack 的博客文章提供了一些很好的示例，展示了可采用的各种归约策略层级，例如线程归约、warp 归约、块归约和集群归约。

因此，只有当批次大小影响归约策略时，我们才会破坏批次不变性。

让我们看看 RMSNorm 的标准并行策略。一般而言，并行算法会受益于尽量减少核心之间的通信。在本文的讨论中，你可以假设我们所说的“核心”指的是 SM。更准确地说，这里重要的性质是：内核启动的线程块数量大于 SM 数量。因此，我们可以从这样一种策略开始：将每个批次元素分配给一个核心，如上图所示。

增大批次大小不会影响归约策略；如果批次大小为 200 时已经能为内核提供足够的并行度，那么批次大小为 2000 时肯定也能提供足够的并行度。

另一方面，减小批次大小可能会带来挑战。由于我们将每个批次元素分配给一个核心，随着批次大小不断减小，最终核心数量会超过批次元素数量，导致部分核心处于空闲状态。

遇到这种情况时，优秀的内核工程师通常会采用上一节提到的某种解决方案（原子加法或拆分归约），以保持良好的并行度，从而维持良好的性能。遗憾的是，这会改变归约策略，导致该内核无法保持批次不变性。

最简单的解决方案就是完全忽略这些情况。这并非完全不合理——较小的批次大小意味着内核很可能无论如何都能快速执行，因此性能下降或许并不会造成灾难性的影响。

如果我们必须优化这一使用场景，一种方法是始终采用一种即使在批次大小非常小时也能提供足够并行度的归约策略。对于较大的批次大小，这种归约策略会产生过多的并行度，但它能让我们在整个尺寸范围内获得尚可（但并非峰值）的性能。

批次不变的矩阵乘法

从本质上讲，你可以把矩阵乘法简单地看作逐点运算之后再进行归约。这样一来，如果我们通过将输出切分成多个分块来并行化矩阵乘法，就能得到一种类似的“数据并行”内核策略，使每次归约都在单个核心内完成。

与 RMSNorm 类似，我们的“批次”维度（M 和 N）也可能变得过小，迫使我们沿归约维度（K）进行拆分。尽管矩阵乘法具有两个“批次”维度，但为了有效利用 tensorcore，我们还需要让每个核心承担更多的“工作”。例如，假设有一个 [1024, K] x [K, 1024] 的矩阵乘法，并采用标准的 [128, 128] 二维分块大小，那么数据并行策略只能将该矩阵乘法拆分到 64 个核心上，不足以让 GPU 达到饱和状态。

在矩阵乘法中沿归约维度进行拆分，称为 Split-K Matmul。与 RMSNorm 一样，使用这种策略会破坏批次不变性。另一种值得关注的矩阵乘法并行策略是 stream-k。stream-k 的有趣之处在于，它甚至比典型的矩阵乘法具有更弱的不变性。如前所述，大多数矩阵乘法库并不具备批次不变性，但至少还具备一种可以称为“批次位置不变性”的性质（即改变元素在批次中的位置不会影响数值结果）。然而，stream-k 连批次位置不变性也不具备！它的核心思想是：针对不同的输出分块，以不同的方式沿 k 维度进行拆分，可以实现更均衡的负载；但利用这一点也会使我们的内核失去批次位置不变性。

矩阵乘法还有一项额外的复杂性——tensor core 指令。对于归约，我们可以简单地一次处理一行；而高效的矩阵乘法内核必须一次处理整个“分块”。

每条 tensor-core 指令（例如 wgmma.mma_async.sync.aligned.m64n128k16）在内部都可能采用不同的归约顺序。使用不同 tensor-core 指令的一个原因，可能是批次大小非常小。例如，如果我们使用一条处理长度为 256 的分块的 tensor-core PTX 指令，但批次大小只有 32，那么几乎所有计算能力都被浪费了！当批次大小为 1 时，最快的内核通常完全不使用 tensor core。

因此，确保矩阵乘法具备批次不变性的最简单方法，是只编译一种内核配置，并将其用于所有形状。尽管这样会损失一部分性能，但在 LLM 推理中通常不会造成灾难性的后果。特别是，只有在 M 和 N 都很小时才最需要 split-k；幸运的是，在我们的场景中，N（即模型维度）通常相当大！

批次不变的注意力机制

在让矩阵乘法具备批次不变性之后，注意力机制又引入了两个额外的复杂因素——这很贴切，因为它包含两次矩阵乘法。

RMSNorm 和矩阵乘法都只沿特征维度进行归约，而现在我们需要同时沿特征维度和序列维度进行归约。

因此，注意力机制必须处理多种会影响序列处理方式的推理优化技术，例如分块预填充、前缀缓存等。

所以，要在 LLM 推理中实现确定性，我们的数值结果必须同时不受以下两点影响：一次处理多少个请求，以及推理引擎如何切分每个请求。

首先来看注意力机制的标准并行策略，该策略最早由 FlashAttention2 提出。与 RMSNorm 和 Matmul 类似，其默认策略也是“数据并行”策略。由于我们沿键张量和值张量进行归约，数据并行策略只能沿查询张量进行并行化。

例如，根据推理引擎的选择，一个序列可能会被分成多个部分处理（如分块预填充），也可能一次性全部处理（如果预填充没有被拆分）。为了实现“批次不变性”，某个给定 token 的归约顺序不能取决于其所在序列中有多少其他 token 正在同时被处理。如果像 vLLM 的 Triton 注意力内核那样，将 KV 缓存中的 K/V 值与当前正在处理的 token 中的 K/V 值分开归约，就无法实现这一点。例如，在处理序列中的第 1000 个查询 token 时，无论 KV 缓存中有 0 个 token（预填充）还是有 999 个 token（解码），归约顺序都必须完全相同。

为了解决这个问题，我们可以在注意力内核本身执行之前更新 KV 缓存和页表，从而确保无论正在处理多少个 token，键和值始终都以一致的方式布局。

加上这一细节，以及上一节提到的所有事项（例如保持一致的分块大小），我们便能够实现批次不变的注意力机制！

然而，这里存在一个严重问题。与矩阵乘法不同，我们在 LLM 推理中遇到的注意力形状往往确实需要使用拆分归约内核，这种内核通常称为 Split-KV 或 FlashDecoding。这是因为，如果不沿归约维度进行并行化，我们就只能沿批次维度、注意力头维度和“查询长度”维度进行并行化。在注意力机制的解码阶段，查询长度非常小，因此除非批次大小非常大，否则我们往往无法让 GPU 达到饱和状态。

遗憾的是，这个问题并不能如此轻易地忽略。