机器之心报道

机器之心编辑部

真正的元凶是缺乏批次不变性。

就在今天，由 OpenAI 前 CTO Mira Murati 成立于今年 2 月的人工智能初创公司 Thinking Machines Lab，发了第一篇文章 ——《克服 LLM 推理中的不确定性》（Defeating Nondeterminism in LLM Inference）。

这篇博客属于 Thinking Machines Lab 新提出的博客栏目 Connectionism，意为「连接主义」。该公司表示：「我们相信，分享才能让科学更好地发展。Connectionism 将涵盖与我们的研究一样广泛的主题：从核函数数值计算到提示工程。Connectionism 这一名称可以追溯到 AI 的早期年代。它曾是 20 世纪 80 年代的一个研究分支，专注于神经网络及其与生物大脑的相似性。」

此外，Thinking Machines Lab 联合创始人、著名技术博主翁荔（Lilian Weng）还在转推中透露了一个消息，Connection Machine，即「连接机」，难道他们的产品要来了？

真是让人期待呢。

地址：https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/

博客主要作者为 Horace He，这位 PyTorch 核心开发者于今年 3 月从 meta 离职，加入了 Thinking Machines。

接下来看博客原文内容。

可复现性（reproducibility）是科学进步的基石。然而，从大语言模型中获得可复现的结果却非常困难。

例如，你可能会发现：即使是向 ChatGPT 提出同一个问题多次，也可能得到不同的回答。这本身并不令人意外，因为语言模型生成结果的过程涉及采样 —— 这个过程会将模型的输出转换为一个概率分布，并以概率方式选择一个 token。

更令人惊讶的是，即使我们将温度参数调到 0（理论上使采样过程变为确定性），大语言模型的 API 在实际中仍然不是确定性的。研究者已经对此有诸多讨论。

即使是在你自己的硬件上，使用开源推理库（如 vLLM 或 SGLang）运行推理，采样过程依然不是确定性的。

为什么大语言模型的推理引擎不是确定性的呢？

一个常见的假设是：浮点运算的非结合性（non-associativity）与并发执行的某种组合会导致不确定性，这取决于哪个并发核心首先完成。我们将这种解释称为「LLM 推理不确定性的『并发 + 浮点』假设」。例如，一篇最近的 arXiv 论文（arXiv:2506.09501）写道：

GPU 中的浮点运算具有非结合性（non-associativity），意味着 (a+b)+c≠a+(b+c)，这是由于精度有限和舍入误差所致。这一特性会直接影响 transformer 架构中注意力得分和 logit 的计算，因为在多线程中进行的并行操作，其执行顺序不同会导致结果差异。

虽然这个假设并不完全错误，但它并没有揭示事情的全貌。

例如，即使在 GPU 上，对相同的数据反复进行相同的矩阵乘法运算，每次的结果也都是每一位都相同的。我们确实在使用浮点数，GPU 也确实具有高度并发性。

那为什么在这个测试中却看不到不确定性呢？

要理解大语言模型推理不确定性的真正原因，我们必须更深入地探究。

不幸的是，甚至连「LLM 推理是确定性」的这一说法的定义都很难明确。或许令人困惑的是，以下这些看似矛盾的说法实际上同时都是真实的：

GPU 上的一些核（kernel）是不确定性的。然而，语言模型在前向传播过程中使用的所有核都是确定性的。此外，像 vLLM 这样的 LLM 推理服务器的前向传播过程，也可以被认为是确定性的。尽管如此，从使用推理服务器的任何用户的角度来看，结果却是不确定性的。

在这篇文章中，我们将解释为什么「并发 + 浮点」假设没有达到目的，揭露 LLM 推理不确定性背后的真正罪魁祸首，并解释如何克服不确定性并在 LLM 推理中获得真正可重复的结果。

原罪：浮点数的非结合性

在讨论不确定性之前，有必要先解释一下为什么存在数值差异。毕竟，我们通常将机器学习模型视为遵循交换律或结合律等结构性规则的数学函数。我们的机器学习库难道不应该提供数学上正确的结果吗？

罪魁祸首是浮点非结合性（floating-point non-associativity）。也就是说，对于浮点数 a、b、c，有：

讽刺的是，正是打破结合律让浮点数变得有用。

浮点数之所以有用，是因为它们允许动态的精度。为了便于解释，我们将使用十进制（而不是二进制），其中浮点数的格式为：尾数 * 10^ 指数。这里还将使用 3 位数字作为尾数，1 位数字作为指数。（注：在计算机科学中，尾数（mantissa，或有效数）是浮点数中用来表示精度的部分，它决定了数字的有效数字位数和精度。）

例如，对于值 3450，我们可以将其精确表示为 3.45*10^3。我们也可以将更小的值（例如 0.486）表示为 4.86*10^-1。这样，浮点数既可以表示非常小的值，也可以表示非常大的值。在科学领域，我们可以说浮点数使我们能够保持有效数的个数恒定。

如果两个浮点数的指数相同，它们的加法运算看起来与整数加法类似。例如：

但是，如果两个浮点数的指数不同，例如 1230 和 23.4，又会发生什么情况呢？理论上，它们的和应该是 1253.4。然而，由于浮点数运算只能保留 3 位有效数字，因此结果会被舍入为 1.25×10³（或 1250）。

表示 1230 需要 3 位有效数字，表示 23.4 也需要 3 位有效数字。但是，这两个数相加的结果（1253.4）却需要 5 位有效数字才能精确表示。因此，我们的浮点数格式必须舍弃最后两位（34）。某种程度上，这相当于我们在相加之前，将原来的 23.4 四舍五入为 20.0。

然而，这样做会导致信息丢失。请注意，只要我们对两个不同阶位（即不同指数）的浮点数进行加法运算，就会发生这种情况。而实际应用中，我们经常需要对不同指数的浮点数进行加法运算。事实上，如果我们能够保证所有浮点数的指数都相同，那么我们完全可以只使用整数！

换句话说，每次以不同顺序相加浮点数时，结果都有可能完全不同。举个极端的例子，对于某个数组，根据加法顺序的不同，其求和结果可能出现 102 种不同的结果。

虽然这是导致输出结果不一致的根本原因，但它并不能直接解释不确定性行为的来源。它也无法帮助我们理解为什么浮点数的加法顺序会改变、这种情况在什么时候发生、以及我们如何避免它。

答案藏在核函数（kernel）的实现方式中。

为什么核函数计算中数字加法顺序并非总是固定的？

如前所述，解释核函数计算中数字加法顺序不一致的一个常见原因是「并发性 + 浮点运算」假设。

该假设认为，如果并发线程的执行顺序是不可预测的，并且累加操作的顺序依赖于并发线程的执行顺序（例如原子加法 /atomic adds），那么最终的累加结果也会变得不可预测。

然而，令人困惑的是，尽管这种现象会导致核函数计算结果的不确定性，但并发机制（以及原子加法）实际上与大型语言模型推理中的不确定性无关！

为了解释真正的罪魁祸首是什么，我们首先需要了解为什么现代 GPU 核函数很少需要使用原子加法。

什么时候需要使用原子加法操作？

GPU 通常会同时在多个核心（即流处理器）上并行运行程序。由于这些核心之间没有内置同步机制，因此如果它们需要相互通信，就会很麻烦。例如，如果所有核心都需要对同一个元素进行累加，就可以使用原子加法（有时也称为 fetch-and-add）。原子加法是不确定性的，结果的累加顺序完全取决于哪个核心先完成计算。

具体来说，假设你要使用 100 个核心对一个包含 100 个元素的向量进行求和（例如 torch.sum ()）。虽然可以并行加载所有 100 个元素，但最终我们必须将结果汇总为一个值。一种实现方法是使用某种原子加法操作，硬件保证所有加法操作都会执行，但并不保证执行顺序。

原子加法操作可以确保每个核心的计算结果都能最终反映在总和中。但是，它并不能保证这些结果的累加顺序。累加顺序完全取决于哪个核心先完成计算，这是一种不确定性行为。

因此，多次执行相同的并行程序可能会产生不同的结果。这通常就是人们所说的不确定性，即，使用完全相同的输入数据执行两次相同的程序，但最终结果却可能不同。这被称为运行间不确定性（run-to-run nondeterminism），例如，运行两次完全相同的 Python 脚本，即使依赖库版本完全相同，结果也可能不同。

虽然并发的原子加法操作会使核函数的执行结果变得不可预测，但对于大多数核函数来说，原子加法并非必需。

事实上，在 LLM 的典型前向传播过程中，通常根本不需要使用原子加法。这可能令人感到意外，因为并行化计算中的归约操作通常可以从原子加法中获益。但实际上，原子加法在大多数情况下并非必需，主要原因有两点。

1. 通常情况下，批处理维度上的并行性已经足够，因此我们无需在归约维度上进行并行化。

2. 随着时间的推移，大多数神经网络库都采用了各种策略，以在不牺牲性能的情况下实现结果的可预测性。

由于上述两个因素，对于绝大多数神经网络操作来说，不使用原子加法几乎不会带来性能损失。

当然，仍然有少数常见操作在不使用原子加法时会遭遇显著的性能下降。例如，PyTorch 中的 scatter_add（即 a [b] += c）。不过，在大语言模型中唯一常用且依赖原子加法的操作，是 FlashAttention 的反向传播（backward）。

然而，LLM 的前向传播过程中并不涉及任何需要原子加法的操作。因此，LLM 的前向过程本质上是运行间确定的（即每次运行结果一致）。

维基百科上写道：一个确定性算法是在给定特定输入的情况下，始终产生相同输出的算法。而在这里，只要输入完全相同（即推理服务器处理的请求完全一致），前向传播就总是会生成完全相同的输出。

然而，前向传播本身是确定性的并不意味着整个系统也是确定性的。比如，如果某个请求的输出依赖于并行用户的请求（例如 batch-norm 这样的操作），那么由于每个请求都无法预知其他并发请求的内容，从单个请求的视角来看，整个 LLM 推理过程就会是不确定性的。

事实证明，我们的请求输出确实依赖于其他并发用户的请求。但这并不是因为跨 batch 泄露了信息，而是因为我们的前向传播过程缺乏批次不变性（batch invariance），这导致同一个请求的输出会受到前向传播中 batch size（batch size）变化的影响。

批次不变性与确定性

为了说明什么是批次不变性，我们可以简化问题，只关注矩阵乘法（matmul）。你可以假设所有的 matmul 实现都是运行间确定的，也就是说，同样的输入，每次运行都会得到相同的结果。

但它们并不是批次不变的。换句话说，当 batch size 发生变化时，batch 中的每个元素可能会得到不同的计算结果。

从数学角度来看，这是一种相当反常的性质。理论上，矩阵乘法在 batch 维度上应当是独立的，batch 中其他元素的存在与否，或 batch 的大小，都不应影响某个具体元素的计算结果。

然而，我们通过实验证据可以发现，现实情况并非如此。

请注意，这里的确定性是指每次运行结果都相同。如果你多次运行该脚本，它会始终返回相同的结果。

但是，如果将非批处理不变的核函数用作更大推理系统的一部分，则整个系统可能变得不确定性。当你向推理端点发送请求时，从用户角度来看，服务器的负载情况是不可预测的。负载决定了核函数的 batch size，从而影响每个请求的最终结果。

如果你把某种核函数不具备不变性的属性（例如：batch size）与该属性本身的不确定性（例如：服务器负载情况）组合在一起，就会得到一个不确定性的系统。

换句话说，几乎所有大语言模型推理端点之所以是不确定的，主要原因就是负载（以及由此决定的 batch size）本身具有不确定性！这种不确定性并非仅限于 GPU，使用 CPU 或 TPU 运行的 LLM 推理端点也会存在同样的问题。因此，如果我们想避免推理服务器中的不确定性，就必须确保核函数对 batch size 具有不变性。

为了理解如何实现这一点，我们首先需要了解为什么核函数默认情况下并不具备批处理不变性。

我们如何使核具有批次不变性？

为了确保 Transformer 模型的实现与 batch size 无关，我们必须确保模型中的每个核心模块都与 batch size 无关。幸运的是，我们可以假设每个逐点运算（pointwise operation）都与 batch size 无关。因此，我们只需要担心涉及的 3 个操作：RMSNorm、矩阵乘法和注意力。

巧合的是，这些操作的难度正好是依次递增的。要想在保持合理性能的同时实现批次不变性，每一种操作都需要一些额外的考量。我们先从 RMSNorm 开始谈起。

RMSNorm

RMSNorm 实现方式：

批次不变性的要求是，无论核函数的 batch size 如何，每个元素的归约顺序都必须保持不变。需要注意的是，这并不意味着我们必须始终使用相同的归约策略。例如，即使我们改变了要进行归约的元素数量，只要归约顺序不变，我们的算法仍然可以满足批处理不变性的要求。

因此，只有当 batch size 影响到归约策略时，我们才会打破批次不变性。

让我们来看一下 RMSNorm 的标准并行化策略。一般来说，并行算法都会从尽量减少核心之间的通信中获益。在这里，为了方便讨论，你可以假设我们所说的核心（cores）就是指 SM（Streaming Multiprocessors，流处理多处理器）。更具体地说，这里重要的性质是：核函数启动的线程块（threadblocks）数量多于 SM 的数量。

基于这一点，一种可行的策略就是：将每个 batch 元素分配给一个核心，就像上图展示的那样。

当我们增加 batch size 时，并不会影响归约策略；如果 batch size = 200 已经能为核函数提供足够的并行性，那么 batch size = 2000 显然也同样能够提供足够的并行性。

另一方面，减小 batch size 也会带来一些挑战。由于我们为每个批次元素分配一个核心，减小 batch size 会导致核心数量大于批次元素数量，从而造成部分核心闲置。遇到这种情况，优秀的核函数工程师会采用前面提到的解决方案之一（原子加法或分段求和），从而保持良好的并行性，进而提升性能。然而，这会改变求和策略，导致该核函数不再具备 batch size 不变的特性。

最简单的解决方案就是直接忽略这些情况。这并不是完全不合理的，因为当 batch size 很小时，核函数通常本来就能很快执行，因此即使出现一些减速，也不会造成灾难性的影响。

如果我们必须优化这种场景，一种方法是：始终使用一种在极小 batch size 下也能提供足够并行度的归约策略。这样的策略会在 batch size 较大时导致过度并行，从而无法达到峰值性能，但它可以让我们在整个 batch size 范围内都获得尚可（虽然不是最佳）的性能表现。

批次不变矩阵乘法

从本质上讲，你可以把矩阵乘法看作是一次逐点运算后接一次归约。那么，如果我们通过将输出划分为小块来并行化矩阵乘法，就能得到一种类似的数据并行核函数策略，使得每一次归约都在单个核心内完成。

与 RMSNorm 类似，矩阵乘法的批次维度（M 和 N）也可能变得过小，迫使我们必须沿归约维度（K）进行拆分。尽管有两个批次维度，矩阵乘法仍然需要每个核心有更多的工作量才能有效利用张量核心。例如，对于一个 [1024, K] x [K, 1024] 的矩阵乘法和一个标准的 [128, 128] 二维 tile 大小，数据并行策略最多只能将其分配到 64 个核心上，这不足以使 GPU 达到饱和。

在矩阵乘法中沿归约维度进行拆分被称为 Split-K 矩阵乘法。与 RMSNorm 的情况一样，使用这种策略会破坏批次不变性。

矩阵乘法还有一个额外的复杂性，即张量核心指令。对于归约操作，我们可以一次只处理一行；但高效的矩阵乘法核函数必须一次性操作一整个 tile。

每条张量核心指令（例如 wgmma.mma_async.sync.aligned.m64n128k16）在内部可能有不同的归约顺序。选择不同张量核心指令的一个原因可能是 batch size 非常小。例如，如果我们使用的张量核心 PTX 指令操作的是一个长度为 256 的 tile，但 batch size 只有 32，那我们几乎浪费了所有的计算资源！当 batch size 为 1 时，最快的核函数通常根本不使用张量核心。

因此，确保矩阵乘法批次不变性的最简单方法是：编译一个固定的核函数配置，并将其用于所有形状的计算。尽管这会损失一些性能，但在 LLM 推理场景下，这种损失通常不是灾难性的。特别是，Split-K 策略在 M 和 N 维度都很小时才最被需要，而幸运的是，在我们的应用场景中，N 维度（即模型维度）通常都相当大！

批次不变性注意力机制

在实现了矩阵乘法的批次不变性之后，注意力机制又引入了两个额外的难题 —— 这也很贴切，因为它正好包含两次矩阵乘法。

1. 与 RMSNorm 和矩阵乘法仅在特征维度上进行归约不同，注意力机制现在需要在特征维度和序列维度上都进行归约。

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

因此，为了在 LLM 推理中实现确定性，我们的数值计算必须对两个因素保持不变：一是单次处理的请求数量，二是每个请求在推理引擎中的切分方式。

我们首先来了解一下注意力机制的标准并行策略，该策略最初由 FlashAttention-2 提出。与 RMSNorm 和矩阵乘法类似，其默认策略是数据并行策略。由于归约是沿着键 / 值（K/V）张量进行的，因此数据并行策略只能沿着查询（Q）张量进行并行化。

例如，根据推理引擎的选择，一个序列可能被分成几个部分处理（如在分块预填充中），也可能一次性处理完毕（如果预填充未被分割）。为了实现批次不变性，对于一个给定的 token，其归约顺序必须独立于其所在序列中同时被处理的其他 token 的数量。

如果你将 KV 缓存中的 K/V 值与当前正在处理的 token 的 K/V 值分开进行归约（就像在 vLLM 的 Triton 注意力核函数中那样），这个目标就无法实现。例如，在处理序列中的第 1000 个查询 token 时，无论 KV 缓存中有 0 个 token（预填充阶段）还是 999 个 token（解码阶段），其归约顺序都必须完全相同。

为解决此问题，我们可以在注意力核函数运行前就更新 KV 缓存和页表，从而确保无论处理多少个 token，我们的键和值始终具有一致的内存布局。

加上这一额外处理（以及前文提到的所有措施，如使用一致的 tile 大小），我们便能实现一个批次不变性的注意力机制！

然而，这里存在一个重要问题。与矩阵乘法不同，LLM 推理中的注意力计算形状通常确实需要一个拆分 - 归约核函数（split-reduction kernel），这类核函数常被称为 Split-KV 或 FlashDecoding。这是因为如果我们不沿着归约维度进行并行，就只能沿着批次维度、头维度和查询长度维度进行并行。

在注意力的解码阶段，查询长度非常小（通常为 1），因此除非 batch size 非常大，否则我们往往无法使 GPU 达到饱和状态。不幸的是，这种情况不像在 RMSNorm 和矩阵乘法中那样容易被忽略。例如，如果你的 KV 缓存非常长，即使只处理一个请求，注意力核函数的计算也可能耗时很长。

此外，常用于注意力的拆分 - 归约策略也给批次不变性带来了挑战。例如，FlashInfer 的平衡调度算法会选择能够使 GPU 所有核心饱和的最大拆分大小，这使得其归约策略并非批次不变的。然而，与 RMSNorm / 矩阵乘法不同，无论 batch size 如何，仅仅选择一个固定的拆分数量是不够的。

相反，为了实现批次不变性，我们必须采用固定拆分大小策略。换言之，我们固定的不是拆分的数量，而是每个拆分块的大小，这样最终会得到一个可变的拆分数量。通过这种方式，我们可以保证无论正在处理多少个 token，我们总是执行完全相同的归约顺序。

实现

我们基于 vLLM，通过利用其 FlexAttention 后端和 torch.Library，提供了一个确定性推理的演示。通过 torch.Library，我们能够以一种非侵入式的方式替换掉大部分相关的 PyTorch 算子。

你可以在 thinking-machines-lab/batch-invariant-ops 找到「批次不变性」核函数库，以及在「确定性」模式下运行的 vLLM 示例。

地址：https://github.com/thinking-machines-lab/batch_invariant_ops

实验

完成结果的不确定性程度如何？

我们使用 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 次结果都变得完全相同。这正是我们期望采样器应有的表现，但若不使用我们的批次不变性核函数，就无法实现确定性结果。

性能

目前，我们还没有投入精力优化批次不变性核函数的性能。不过，我们还是进行了一些实验来验证其性能是否仍在可用范围内。

我们搭建了一个配备单块 GPU 的 API 服务器，运行 Qwen-3-8B 模型，并请求生成 1000 个序列，输出长度控制在 90 到 110 个 token 之间。

性能下降的主要原因在于 vLLM 中的 FlexAttention 集成尚未经过深度优化。尽管如此，我们看到其性能并未出现灾难性下降。

真正的在策略强化学习

正如研究人员所指出的，训练和推理之间的数值差异会隐式地将我们的在策略强化学习（on-policy RL）转变为离策略强化学习（off-policy RL）。

当然，如果我们甚至无法从两次相同的推理请求中获得每一位都相同的结果，那么在训练和推理之间获得每一位都相同的结果也是不可能的。因此，确定性推理使我们能够修改训练堆栈，从而在采样和训练之间获得每一位都相同的结果，最终实现真正的在策略强化学习。

我们在 Bigmath 上，使用 RLVR 设置进行了实验，其中强化学习策略由 Qwen 2.5-VL instruct 8B 模型初始化，最大 rollout 长度为 4096。

如果我们不使用离策略校正（即重要度加权）进行训练，我们的奖励会在训练中途崩溃；而添加离策略校正项则可以使训练顺利进行。但是，如果我们在采样器和训练器之间实现了每一位都相同的结果，我们就完全处于在策略状态（即 KL 散度为 0），同样可以顺利地进行训练。

我们还可以绘制采样器和训练器之间对数概率的 KL 散度，其中所有 3 次运行都表现出显著不同的行为。在使用重要度加权运行时，KL 散度保持在 0.001 左右，并伴有偶尔的峰值。然而，在不使用重要度加权的情况下运行，最终会导致 KL 散度在大约与奖励崩溃同一时间出现峰值。当然，在运行「真正的在策略强化学习」时，我们的 KL 散度始终保持为 0，这表明训练策略和采样策略之间不存在任何差异。

总结

现代软件系统往往由多层抽象构成。在机器学习中，当我们遇到不确定性和一些微妙的数值差异时，人们往往会倾向于视而不见。

毕竟，我们的系统本来就是「概率性的」，再多一点不确定性又有何妨？单元测试挂掉时，把 atol/rtol 调大点有什么问题？训练器和采样器之间的对数概率差异，应该不是真正的 bug 吧？

我们拒绝这种消极心态。只要稍微多做一些努力，我们就能理解不确定性的根源，甚至真正解决它们！

我们希望这篇博文能为社区提供一套可靠的思路，帮助大家在推理系统中应对不确定性，并激励更多人深入理解自己的系统。