Thinking Machines Lab 解析 LLM 推理不确定性成因并开源 batch-invariant 内核

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

比如，你可能会发现，向 ChatGPT 反复提同一个问题，得到的回答却不一样。这本身并不奇怪，因为从语言模型得到结果要经过“采样”：这个过程把模型的输出转成概率分布，再按概率挑出一个 token。

更让人意外的是，即便把温度调到 0这意味着 LLM 总是选择概率最高的 token，也就是所谓的贪心采样。（理论上采样因此变成确定性的），LLM API 在实际使用中仍然不是确定性的（参见以往的讨论这里、这里、这里）。就算用 vLLM、SGLang 这类开源推理库在自己的硬件上跑推理，采样依然不是确定性的（参见这里、这里）。

那么，LLM 推理引擎为什么不是确定性的？一种常见的假设是：浮点运算不满足结合律，再加上并发执行，最终结果取决于哪个并发核心先完成，于是产生了不确定性。我们把这个假设称为 LLM 推理不确定性的“并发 + 浮点”假说。例如，最近的一篇 arXiv 预印本写道：

GPU 中的浮点运算不满足结合律，即 (a+b)+c≠a+(b+c)(a+b)+c=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 上有些 kernel 是非确定性的。
但语言模型前向传播中用到的所有 kernel 都是确定性的。
而且，LLM 推理服务器（比如 vLLM）的前向传播也可以说是确定性的。
然而，从任何使用推理服务器的人的角度看，结果却是非确定性的。

本文会解释为什么“并发 + 浮点”这个假说并不成立，揭开 LLM 推理非确定性的真正元凶，并说明如何消除非确定性，在 LLM 推理中获得真正可复现的结果。

原罪：浮点运算不满足结合律

在讨论非确定性之前，有必要先解释为什么会出现数值差异。毕竟我们通常把机器学习模型看作遵循交换律、结合律等结构规则的数学函数，机器学习库难道不应该给出一个“数学上正确”的结果吗？

罪魁祸首是浮点运算不满足结合律。也就是说，对浮点数而言：

(a+b)+c≠a+(b+c)(a+b)+c=a+(b+c)

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

讽刺的是，浮点数之所以有用，恰恰是因为它打破了结合律。

浮点数之所以有用，是因为它允许“动态”的精度。为便于说明，这里用十进制（而不是二进制），浮点数的格式为尾数∗10指数尾数∗10指数。同时，尾数用 3 位数字，指数用 1 位数字。

例如，对于数值 3450，我们可以把它精确表示为 3.45∗1033.45∗103。像 0.486 这样小得多的数值，也可以表示为 4.86∗10−14.86∗10−1。这样一来，浮点数就能同时表示非常小和非常大的数值。在科学领域，我们也可以说，浮点数让我们能够保持固定的“有效数字”位数。

如果把两个指数相同的浮点数相加，过程和整数加法很像。例如，123（1.23∗1021.23∗102）+ 456（4.56∗1024.56∗102）的结果是 579（5.79∗1025.79∗102）。

但如果把两个指数不同的浮点数相加，比如 1230 和 23.4，会发生什么？这种情况下，精确结果是 1253.4。然而，我们一次只能保留 3 位精度。因此，浮点加法会丢弃最后 2 位数字，得到 1.25∗1031.25∗103（也就是 1250）。

1.23 × 10²

3.45 × 10¹

1.575 × 10²

精确值：1575

表示 1230 需要 3 位精度，表示 23.4 也需要 3 位精度。但把这两个数相加，得到的数需要 5 位精度才能表示（1253.4）。于是，我们的浮点格式只能把末尾的 34 丢掉。从某种意义上说，我们实际上是在相加之前，先把原来的 23.4 舍入成了 20.0。

然而，到这一步，信息已经被破坏了。注意，每当我们把两个“尺度”不同（即指数不同）的浮点数相加时，都可能发生这种情况。而指数不同的浮点数相加，随时都在发生。事实上，如果我们能保证永远不需要不同的指数，那直接用整数就行了！

换句话说，我们每次以不同的顺序把浮点数相加，都可能得到完全不同的结果。举个极端的例子：对这个数组求和，根据相加顺序的不同，可能得到 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 实现在反向传播中做了额外的重计算，从而避免原子加，但代价是多花 40% 的 FLOPs！

不过，LLM 的前向传播不涉及任何需要原子加的运算。因此，LLM 的前向传播实际上是“逐次运行确定”的。

从推理服务器的角度看，它确实是确定的。给定完全相同的用户请求，它总是给出相同的确定性输出。

维基百科写道：“确定性算法是指，给定特定输入，总是产生相同输出的算法。”在这个例子里，给定完全相同的输入（即推理服务器正在处理的完全相同的请求），前向传播总是产生完全相同的输出。

然而，前向传播本身“确定”，并不足以保证包含它的系统是确定的。例如，如果我们的请求输出依赖于并行的用户请求（比如 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 推理是“不确定”的。

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

换句话说，几乎所有 LLM 推理端点都不确定，首要原因就在于负载（以及由此决定的批次大小）会不确定地变化！这种不确定性并非 GPU 独有——由 CPU 或 TPU 提供服务的 LLM 推理端点同样存在这一不确定性来源。

所以，如果我们想避免推理服务器中的不确定性，就必须让内核实现批次不变。为了理解如何做到这一点，我们先来看看内核一开始为什么不具备批次不变性。

如何让内核实现批次不变？

要让 transformer 实现批次不变，我们必须让每个内核都批次不变。幸运的是，我们可以假定所有逐点操作都是批次不变的。虽然对 PyTorch 中的所有内核来说确实如此，但这并非天然成立。例如，CPU 上有些内核实现会对数组的某些部分使用向量化内建指令，对其他部分使用非向量化内建指令，而这些内建指令的数值结果未必总是逐位相同。因此，我们只需要关注涉及归约的 3 个操作——RMSNorm、矩阵乘法和注意力。与并行相关的归约不在本文讨论范围内，但同样的原则适用。一个可能有用的小知识是：在 Blackwell 以及搭配 CUDA 12.8+ 的 Hopper 上，NVLink-Sharp 的交换机内归约是确定性的。和许多事情一样，这一信息可以在 NCCL 的 github issues 中找到

方便的是，这几个操作的难度也恰好按升序排列。每一个都需要一些额外考量，才能在合理性能下实现批次不变。我们先从 RMSNorm 说起。

批次不变的 RMSNorm

**数据并行 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

批次不变性的要求是：无论内核的批次大小如何，每个元素的归约顺序必须固定。注意，这并不意味着我们必须始终使用相同的归约策略。例如，即使我们改变归约的元素数量，只要归约策略的变化不影响顺序，我们仍然可以保持批次不变性。The Quack 博客文章中有一些很好的例子，展示了你可以采用的各种归约策略的层次结构（例如线程归约、warp 归约、block 归约、cluster 归约）。

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

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

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

**面向更大批次的数据并行 RMSNorm** 将数据并行策略扩展到更大的批次相当直接——不是让每个核心处理一行，而是让每个核心依次处理不同的行。这*保持了批次不变性*，因为每个批次元素的归约策略保持不变。

另一方面，减小批次大小可能会带来挑战。因为我们把每个批次元素分配给一个核心，减小批次大小最终会导致核心数量多于批次元素，使一些核心空闲。

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

**拆分归约 RMSNorm** 如果批次大小很小，我们的数据并行策略可能不再有足够的并行性来让核心满载。在这种情况下，将归约“拆分”到多个核心上可能更高效，让我们充分利用 GPU。然而，这会失去批次不变性，因为我们不再以相同的顺序归约每个元素。

最简单的解决方案就是完全忽略这些情况。这并非完全不合理——批次大小小意味着内核很可能执行得很快，因此性能下降可能并不严重。

如果我们确实必须针对这个用例做优化，一种思路是始终采用一种归约策略，让它在极小的批大小下也能有足够的并行度。这种归约策略在批大小较大时会产生过剩的并行度，但能让我们在整个批大小范围内都获得尚可（而非峰值）的性能。

批不变矩阵乘法

**Data Parallel Matmul** Similar to RMSNorm, the standard parallelism strategy for matmuls is a “data-parallel” strategy, keeping the entire reduction in one core. It is most straightforward to think about splitting the output tensor into 2D tiles and assigning each tile to a different core. Each core then computes the dot products that belong to that tile, once again performing the entire reduction within one core.
与 RMSNorm 不同，算术强度和利用 tensor core 方面的额外约束，迫使我们在高效的 matmul 内核中拆分 2D 瓦片，而不是拆分单个输出元素。

从本质上讲，矩阵乘法可以看作先做逐点运算、再做归约。那么，如果我们把输出切分成瓦片来并行化矩阵乘法，就得到了一种类似的“数据并行”内核策略，让每次归约都留在同一个核心内完成。

同样与 RMSNorm 类似，我们的“批”维度（M 和 N）也可能变得太小，从而被迫沿归约维度（K）拆分。尽管有两个“批”维度，matmul 仍要求每个核心承担多得多的“工作量”，才能有效利用 tensor core。例如，对于一个 [1024, K] x [K, 1024] 的 matmul 和标准的 [128, 128] 二维瓦片大小，数据并行策略只能把这个 matmul 拆分到 64 个核心上，不足以让 GPU 满载。

在 matmul 中沿归约维度拆分被称为 Split-K Matmul。和 RMSNorm 一样，使用这种策略会破坏批不变性。 matmul 还有一种有趣的并行策略叫 stream-k。stream-k 的有趣之处在于，它的不变性比典型 matmul 还要少。如前所述，大多数 matmul 库都不是批不变的，但至少可以称得上是批位置不变的（也就是说，改变元素在批内的位置不会影响数值结果）。然而，stream-k 连批位置不变性都不具备！它的核心洞见是：针对不同的输出瓦片，用不同的方式沿 k 拆分，可以获得更干净的负载均衡；但利用这一点，我们的内核也就不再是批位置不变的了。

**Split-K Matmul** 如果批维度相当小，我们可能没有足够的并行度，就需要 split-k matmul。在这个例子中，我们把每次归约拆分到两个核心上，各自单独累加，最后再合并结果。不过，把每次归约拆分到两个核心上，让我们仍然能用上八个核心。

matmul 还有一层额外的复杂性——tensor core 指令。归约可以一次只处理一行，而高效的矩阵乘法内核必须一次处理整个“瓦片”。

每条 tensor core 指令（比如 wgmma.mma_async.sync.aligned.m64n128k16）内部的归约顺序可能各不相同。需要换用不同 tensor core 指令的一个原因，可能是批大小非常小。例如，如果我们使用一条处理长度为 256 的瓦片的 tensor core PTX 指令，而批大小只有 32，那几乎全部算力都被浪费了！在批大小为 1 时，最快的内核通常根本不用 tensor core。

**填充式 Tensor-Core 指令** 如果批次规模太小，就可能出现连一个 2D tile 都放不进输出的情况。这时最高效的做法是改用更小的 tensor-core 指令，或者干脆不用 tensor-core！但这两种做法都会让我们的 kernel 无法做到批次不变。

所以，要保证 matmul 的批次不变性，最简单的办法是只编译一种 kernel 配置，所有形状都用它。这样会损失一些性能，但在 LLM 推理中通常不算灾难。尤其是 split-k 最需要用在 M 和 N 都很小的时候，而幸运的是，我们这里 N（即模型维度）通常相当大！

虽然实现了批次不变，但相比 cuBLAS 我们只损失了约 20% 的性能。注意这也不是优化过的 Triton kernel（比如没用 TMA）。不过，性能曲线中的一些形态恰好说明了批次不变的要求在哪里损失了性能。首先，在极小的批次规模下，由于指令过大、并行度不足，性能损失明显。其次，随着批次规模增大会出现“拼图”状的波动，这是量化效应（tile 和 wave 两方面）造成的，通常靠调整 tile 尺寸来缓解。关于这些量化效应，可以看这里。

批次不变的注意力机制

**FlashAttention2 策略** 我们沿 Q 并行，同时沿 K/V 做归约。这意味着整个归约可以留在单个核内完成，因此它也是一种数据并行策略。

在 matmul 实现批次不变之后，注意力机制又带来两个额外的麻烦——说来也巧，因为它包含两个 matmul。

与 RMSNorm 和 matmul 只沿特征维度归约不同，这里我们要同时沿特征维度和序列维度归约。
正因如此，注意力机制必须应对各种影响序列处理方式的推理优化（分块预填充、前缀缓存等）。

因此，要在 LLM 推理中实现确定性，我们的数值计算必须对两件事都不变：同时处理多少个请求，以及每个请求在推理引擎中被如何切分。

我们先梳理注意力的标准并行策略，它最早在 FlashAttention2 中提出。和 RMSNorm、Matmul 类似，默认策略是“数据并行”。由于我们沿 key/value 张量归约，数据并行策略只能沿 query 张量并行。

例如，取决于推理引擎的选择，一个序列可能被分成几段处理（比如分块预填充），也可能一次性处理完（如果预填充没有拆分）。要实现“批次不变”，某个 token 的归约顺序不能取决于它的序列中还有多少其他 token 正在被同时处理。如果你把 KV 缓存中的 K/V 值和当前正在处理的 token 的 K/V 值分开归约（就像 vLLM 的 Triton attention kernel 那样），就无法做到这一点。例如，处理序列中第 1000 个 query token 时，无论 KV 缓存里有 0 个 token（预填充）还是 999 个 token（解码），归约顺序都必须完全一致。

**FlashAttention with a KV Cache** The reason why explicitly handling the KV cache separately from the current KV values breaks batch invariance is a bit subtle and is related to “boundary conditions”. In particular, imagine your block size is 32 but we currently have 80 elements in our KV cache. We then compute an additional 48 elements that aren’t cached. In this case, we need three blocks (two full and one masked) to compute “P cache” and another two blocks (one full and one masked) to compute “P”. This is therefore five total blocks to compute our reduction when we only have four total blocks (i.e. 128) of elements to compute, which will definitely change our reduction order.
举个例子，假设 KV Cache 里没有任何元素、而我们要一次性处理 128 个元素，这两种情况下必须得到完全相同的数值结果，才能保证注意力计算的“批次不变性”。

要解决这个问题，我们只需在注意力核函数执行之前先更新 KV cache 和页表，这样无论当前处理多少个 token，键和值的布局都始终一致。

加上这个细节（以及上一节提到的所有措施，比如固定 tile 大小），我们就能实现批次不变的注意力计算了！

但这里有个大问题。与矩阵乘法不同，LLM 推理中遇到的注意力形状往往确实需要拆分归约核函数，通常称为 Split-KV 或 FlashDecoding。原因在于，如果不沿归约维度并行，我们就只能沿批次维度、头维度和“查询长度”维度并行。在注意力的解码阶段，查询长度非常小，所以除非批次规模非常大，否则往往无法把 GPU 跑满。

可惜的是，这个情况不像 RMSNorm 和矩阵乘法那样容易忽略。比如 KV cache 很长时，哪怕只处理一个请求，注意力核函数也可能耗时很久。

**Fixed # Split-KV Strategy (i.e. FlashDecode)** If our query length becomes very small (like it does during decoding), we may end up in a situation where there is very little parallelism in our kernel at all. In these cases, we’ll need to once again split along the reduction dimension—the KV dimension this time. The typical strategy for how to split along the KV dimension is to figure out how much parallelism we need and then divide the KV dimension evenly. For example, if our KV length was 1000 and we needed 4 splits, each core would handle 250 elements.
这同样会破坏批次不变性，因为我们具体的归约策略取决于当前请求要处理多少个查询 token。

此外，注意力常用的拆分归约策略本身也给批次不变性带来挑战。例如 FlashInfer 的“均衡调度算法”会选择能跑满所有 GPU 核心的最大拆分尺寸，这就使归约策略不具备“批次不变性”。但与 RMSNorm/矩阵乘法不同的是，单纯固定拆分数量、不管批次大小，并不够。

要实现批次不变性，我们必须采用“固定拆分尺寸”的策略。也就是说，不固定拆分数量，而是固定每个拆分的尺寸，拆分数量则随之变化。这样无论处理多少 token，我们都能保证归约顺序完全一致。这需要对 FlexAttention 做一些内部改动，尚未包含在我们的代码发布中。我们会在近期向上游提交！

**Fixed Size Split-KV Strategy** The only difference between this strategy and the previous strategy is that our splits are now “fixed size”. For example, if our KV length was 1000, instead of splitting it into four even length 250 splits, we would split it into three fixed-size length 256 splits and one length 232 split.
这样一来，归约策略不再取决于一次处理多少个查询 token，我们就能保持批次不变性！

实现

我们在 vLLM 之上提供了一个确定性推理的演示，借助它的 FlexAttention 后端以及 torch.Library。通过 torch.Library，我们可以用非侵入的方式替换掉大部分相关的 PyTorch 算子。“批次不变”核函数库见 thinking-machines-lab/batch-invariant-ops，另外还有以“确定性”模式运行 vLLM 的示例。

实验

补全结果到底有多不确定？

我们用 Qwen/Qwen3-235B-A22B-Instruct-2507，在温度 0 下以“Tell me about Richard Feynman”为提示词（非思考模式）采样 1000 次补全，每次生成 1000 个 token。令人意外的是，我们得到了 80 种不同的补全，其中最常见的一种出现了 78 次。

对比各条补全结果出现差异的位置，会发现前 102 个 token 的补全内容其实完全一致！第一次出现分歧是在第 103 个 token。所有补全都生成了“Feynman was born on May 11, 1918, in”这一序列，但其中 992 条补全接着生成“Queens, New York”，另外 8 条则生成“New York City”。

另一方面，启用批不变内核后，我们的 1000 条补全全部相同。从数学上讲，这正是我们对采样器的预期，但如果不使用批不变内核，就无法得到确定性的结果。

性能

我们并没有投入太多精力去优化批不变内核的性能。不过，还是做几组实验，确认性能仍然可用。

我们搭一个 API 服务器，用一块 GPU 跑 Qwen-3-8B，请求 1000 条输出长度在 90 到 110 之间的序列。

配置	耗时（秒）
vLLM 默认	26
未优化的确定性 vLLM	55
+ 改进的注意力内核	42

速度变慢很大程度上是因为 vLLM 里的 FlexAttention 集成还没有经过充分优化。尽管如此，性能也谈不上灾难性。

真正的同策略强化学习

正如研究者已经指出的，训练与推理之间的数值差异，实际上把我们的同策略强化学习变成了异策略强化学习。

当然，如果连两次完全相同的推理请求都无法得到逐位一致的结果，那训练和推理之间就更不可能做到逐位一致了。而确定性推理让我们可以进一步改造训练栈，使采样与训练之间也得到逐位一致的结果，从而实现真正的同策略强化学习。

我们在 Bigmath 上做 RLVR 实验，强化学习策略从 Qwen 2.5-VL instruct 8B 初始化，最大 rollout 长度为 4096。

如果不做异策略校正（即重要性加权），奖励会在训练中途崩溃；而加入异策略校正项后，训练就能平稳进行。不过，如果采样器和训练器之间能做到逐位一致，我们就完全处于同策略（即 KL 散度为 0），同样可以平稳训练。

我们还可以画出采样器与训练器之间 logprobs 的 KL 散度，三次运行的表现明显不同。使用重要性加权时，KL 散度维持在 0.001 左右，偶尔出现尖峰。而不使用重要性加权时，KL 散度最终会在奖励崩溃的同一时间点附近出现尖峰。当然，在“真正的同策略强化学习”下，KL 散度始终保持在 0，说明训练策略与采样策略之间没有任何散度。

注意，不使用重要性加权的那次运行在第 318 步附近出现了明显的损失尖峰，同时 logprobs 的 KL 散度也相应出现尖峰。而无论是使用异策略校正，还是采用“真正的同策略”，都能让强化学习平稳继续。标为“True On-Policy”的蓝线不是 bug——它就是一条位于 0 的平直线。

结论

现代软件系统包含许多层抽象。在机器学习中，当我们遇到不确定性和细微的数值差异时，往往很容易想把它们掩盖过去。毕竟，我们的系统本来就是“概率性”的，再多一点不确定性又有什么关系呢？把失败的单元测试里的 atol/rtol 调大一点又有什么问题呢？训练器和采样器之间 logprobs 的差异大概不是真正的 bug，对吧？

我们拒绝这种消极态度。只要花一点功夫，我们确实能够理解不确定性的根本原因，甚至解决它们！我们希望这篇博客文章能让社区扎实地理解如何解决推理系统中的不确定性问题，并激励其他人去全面理解自己的系统。

引用

请按如下方式引用本工作：

He, Horace and Thinking Machines Lab, "Defeating Nondeterminism in LLM Inference", 
Thinking Machines Lab: Connectionism, Sep 2025.

或使用 BibTeX 引用：

@article{he2025nondeterminism,
  author = {Horace He and Thinking Machines Lab},
  title = {Defeating Nondeterminism in LLM Inference},
  journal = {Thinking Machines Lab: Connectionism},
  year = {2025},
  note = {https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/},
  doi = {10.64434/tml.20250910}
}