成立半年估值$120亿，Thinking Machines Lab首次出击，为何发了篇技术博客？

文｜四木洋子

在 AI 巨头们争相比拼声量、疯狂抢夺人才的浮躁时代，有一家公司选择了截然不同的路径 —— 它就是 Thinking Machines。

这家由 OpenAI 前 CTO 米拉·穆拉蒂（Mira Murati）创立的新兴 AI 公司，成立至今还不到半年，却已经汇聚了一众来自 OpenAI、Meta 等顶级 AI 机构的精英人才。

更令人瞩目的是，今年夏天，Thinking Machines 在尚未发布任何产品的情况下，就完成了20亿美元的种子轮融资，公司估值飙升至120亿美元。

在数月的低调后，这家备受瞩目的AI新星终于有了新动作——上周，它悄然推出了自己的技术博客专栏，并发布了首篇技术文章 "Defeating Nondeterminism in LLM Inference"（《战胜大模型推理中的不确定性》）。

*原文：https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference

这篇深度技术文章直指一个困扰业界已久的核心难题：大语言模型前向推理过程中的不确定性问题。

这个问题从 LLM 诞生之初就如影随形，但业界一直苦于找不到理想的解决方案。

而 Thinking Machines 的这篇文章不仅从算法、系统和硬件三个维度深入剖析了问题的根源所在，更难能可贵的是，他们还提出了切实可行的解决方案——这或许正是这家"神秘"公司以此作为首秀的原因所在。

经过实验验证，Thinking Machines 提出的方案体现出双重突破：

• 一方面，它让大语言模型能够提供高度可复现的前向推理结果，解决了长期困扰研究者的一致性问题；

• 另一方面，该方案从根本上改善了强化学习的训练动态，成功实现了基于大语言模型的真正同策略探索与利用平衡——这一直是这个领域亟待攻克的核心难题。

这些突破性成果印证了 Thinking Machines 团队在 AI 前沿研究中所具备的深厚洞察力和卓越的系统性研究能力。

我们把本文的阅读心得，分为内容总结、思考点评和原文翻译三部分，供大家按需了解：

内容总结

为什么问大模型相同的问题，会得到不同的答案？

这个问题虽然有时会有点恼人，但并不难理解——因为语言模型中的结果涉及“采样”，这本身就是一个"随机选词&输出"的过程。

如果说随机采样约等于抽签，那么采样温度（temperature）则直接决定了抽签的权重分布。

但现在的问题在于，即便将采样温度调整到 0，我们仍然不能保证模型会输出确定的答案。

*temperature = 0 意味着大模型预测过程总是选择最高概率的词元，这被称为贪婪采样，其结果总是确定的。

这就有点令人费解了。

Thinking Machines 的这篇文章从一个我们都不陌生的现象切入——同样的问题，大语言模型有时会给出不同的答案。这看似随机的表现，实际上揭示了大模型推理过程中一个深层次的技术问题：推理的"不确定性"（nondeterminism）。

文章的核心发现令人意外：这种不确定性的根源，竟然来自于推理批次不变性（batch invariance）的缺失。

什么意思呢？让我们从 GPU 的工作机制说起。为了榨干硬件的每一分算力，GPU 会根据当前的批次大小（batch size）——也就是服务器同时处理的请求数量——动态调整自己的并行策略和计算单元的调度方式。批次大了，就用这种策略；批次小了，就换那种策略。

问题就出在这里。

由于浮点数运算本身的特性——它并不严格满足数学上的加法结合律，不同的计算顺序会产生微小的数值偏差。这些看似微不足道的差异，在 Transformer 层层叠叠的前向传播过程中被逐步放大，最终导致了一个令人困扰的结果：相同的输入，在不同的批次环境下，竟然会产生不同的输出。

简而言之，GPU 为了优化效率而采用的灵活调度策略，意外地成为了大模型推理不确定性的幕后推手。一个追求极致性能的设计，却在无意中为模型的"善变"埋下了伏笔。

那么，Thinking Machines 是怎么解决的呢？简单概括，让模型的计算对批次大小敏感因素“免疫”。

在 GPU 计算中，每个操作都通过专门的计算内核（kernel）来执行，这就要求所有内核都必须具备批次不变性——即无论输入批次如何变化，计算逻辑都应保持一致。

然而，在 Transformer 模型的实际运行中，某些操作天然地会因批次规模的改变而产生计算行为上的差异，这类操作主要集中在各种归约计算（reduction）环节。认识到这一挑战，Thinking Machines针对性地对三个核心的归约操作 —— RMSNorm、矩阵乘法以及注意力机制 —— 进行了专门的优化设计。

• RMSNorm：让每一个样本都单独在一个内核上完成 RMSNorm 的计算。这样一来，无论批次里有多少样本，都不需要跨样本通信，完全在自己的内核里完成，得到的结果与单独算该样本时是一致的 。

• 矩阵乘法：简单起见，可以把矩阵乘看作“逐元素相乘再累加”的过程。实现批次不变性的原则同样是让每个输出元素的累加都在单一内核中按固定顺序完成。具体方法是将输出矩阵划分成很多小块（tile），每个内核负责自己那个块的所有计算，并把该块内部需要累加的部分串行做完。这样每个输出块的加法顺序是固定的，不受别的核心影响。

• 注意力机制：注意力计算涉及更复杂的步骤（比如 KV 的点积 和 Softmax 归一化等），而且还牵涉序列长度（token 数量）维度的变化。解决方案是固定归约块的大小。简单理解：原先注意力在解码时为了适应不同 token 数，可能动态地选择需要多少块来并行计算；现在改为预先固定每一块的大小，这样无论总共多少 token，就算需要的块数变多，每个块内部的计算顺序和规模都是一致的。

通过这三种归约计算方法的优化，Thinking Machines 实现了 Transformer 模型前向计算的批次确定性。

有了确定性之后，除了能够实现可复现的模型推理结果，还能显著改善 RL 训练效果。

RL 的基本原理是通过奖励机制来强化模型的正确输出。然而，如果每次答案都有细微差别，训练数据就会变得很「嘈杂」，而更一致的响应能让整个 RL 训练过程变得「丝滑」很多。

在 RL 领域，On-policy（同策略）和 Off-policy（异策略）是两种核心训练范式，它们的区别在于学习时使用的"行为策略"与"目标策略"是否为同一个策略。目前，很多研究工作都在强调同策略（On-policy）优化的重要性。

但这里出现了一个有趣的问题：GPU 的不确定性既可能影响采样时的策略表现（生成不同的数据），也可能影响训练时的策略表现（产生不同的梯度计算结果）。

这就导致了一个核心矛盾——明明是同一个策略，但在采样和训练阶段跑出来的结果却不一致。结果，原本设计为同策略的训练被迫变成了异策略（Off-policy）训练，这种现象可以称为"伪同策略 RL"（Faked On-Policy RL）。

解决方案是通过确定性推理（通过 batch-invariant kernel 等方式），确保推理和训练逐位一致，保证真正的同策略训练。

思考点评

大咖云集的 Thinking Machines 为啥选择这项工作作为自己的首次亮相？

答案很简单：这个问题实在太重要了，而能够针对这一问题提出系统性解决方案的团队却寥寥无几。

Faked On-policy RL for LLM，这是当前世界范围内普遍存在的痛点问题。

令人意外的是，即便是那些在前沿的 AI 顶会上发表、明确宣称"我们采用同策略强化学习"的研究工作，实际情况往往是：研究者们在算法层面精心打磨，但承载算法实现的底层框架和算子等系统组件，却在无形中背离了 On-Policy RL 的初衷和愿景。

而这一根本性问题，恰恰是目前开源社区（如 veRL、MS-Swift 等 RL 训练框架的提供方）难以攻克的 ——原因在于维护人员往往术业有专攻：擅长框架设计的未必精通算子开发，而算子专家又不一定能驾驭整体框架架构。这种技能壁垒的存在，使得端到端的解决方案变得格外珍贵。

目前我们看到的相似研究风格的博客包括但不限于以下几个：

• 如 Thinking Machines 这篇文章中所引用的，微软高剑峰组在前两个月发现了目前 RL for LLM 训练中普遍存在的这一问题。虽然他们在《Your Efficient RL Framework Secretly Brings You Off-Policy RL Training》中提出了缓解方案，但坦率地说，这种方式仍显得不够优雅且不够彻底。

• 同样在算法与系统（算子）层面寻求改进的，还有最近以兼职身份加入 Thinking Machines 的 Songlin Yang 提出的 DeltaNet 系列 也是类似的风格。

• 其他更多的博客文章则侧重于单点优化和改进，如系统（算子）层面 Tri Dao 的 FlashAttention 系列，以及纯算法层面的典型如 John Schulman 的《Approximating KL Divergence》。

而 Thinking Machines 的做法，实际上是拿软件（算法）—— 系统—— 硬件这三者串联所产生的推理不确定性问题，进行了一次由表及里的深度剖析，系统性地梳理推理过程中不确定性的根本来源。

他们的研究路径颇有章法：

从「数值计算」的底层细节，到「系统（内核）原理」的中层机制，再到算法层面的优化改进，层层递进，向外界展示了那些在一般研究人员眼中"平平无奇"的技术细节，是如何最终影响到终端模型训练效果的。

更重要的是，在解决了高剑峰等人发现的「现在的 RL for LLM 都是伪同策略 RL」这一关键问题之后，才有可能在此基础上构建出真正的满血版 RL-ed LLMs（经过强化学习训练的语言模型）。

而这样的模型，又将为其他更宏大的目标提供坚实支撑——比如构建多模态通用 Agent，或是提供面向客户的优质 RL 解决方案等。

在文章的结尾部分，Thinking Machines 的研究人员特别强调：

不能忽视推理过程中的不确定性，只有抓住并深入分析这些不确定性，找到切实可行的解决方案，才能最终实现真正的同策略强化学习，进而把 RL for LLM 这件事做得更加扎实。

对其他投身大模型研究的人来说，Thinking Machines 这次的工作传递出一个清晰的信号：

想要继续提升模型表现，既需要自顶向下的宏观思考，也需要自底向上的细致打磨，端到端解决问题的系统思维和扎实的动手能力缺一不可——「宏观规划 + 微观操作」的组合拳必须打得漂亮。

以下为原文翻译：

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

例如，你可能会发现，多次向 ChatGPT 提出相同的问题会得到不同的结果。这本身并不令人惊讶，因为从语言模型中获取结果涉及“采样”，这是一个将语言模型输出转换为概率分布并概率性地选择一个标记的过程。

更令人惊讶的可能是，即使我们将温度调整到 0（这意味着 LLM 总是选择最高概率的标记，这被称为贪婪采样），LLM API 在实践中仍然不是确定性的。即使在你自己的硬件上使用 vLLM 或 SGLang 等开源推理库运行推理，采样仍然不是确定性的。

但为什么 LLM 推理引擎不是100%确定性的呢？一个常见的假设是，浮点非关联性和并发执行的某种组合导致了基于哪个并发核心首先完成的不确定性。我们将此称为 LLM 推理不确定性的“并发 + 浮点”假设。例如，最近的一篇 arXiv 文章写道：

“GPU 中的浮点运算表现出非关联性，这意味着 (a + b) + c ≠ a + (b + c)，这是由于有限精度和舍入误差造成的。此属性直接影响 Transformer 架构中注意力分数和词表输出分数（logits）的计算，其中跨多个线程的并行操作可能会根据执行顺序产生不同的结果。”

你还可以在其他地方找到“并发 + 浮点”假设的重复，例如：“存在速度权衡，为了使端点快速，使用了 GPU，它进行并行不确定性计算。任何现代 GPU 神经网络计算都将受制于这些”，或“因为 GPU 是高度并行化的，所以每次执行加法或乘法的顺序可能不同，这可能会导致输出中的微小差异”。

虽然这个假设并非完全错误，但它并未揭示全貌。例如，即使在 GPU 上，反复对相同的数据执行相同的矩阵乘法运算，每次得到的结果在二进制层面都是完全相同的。我们确实在用浮点数运算，而且 GPU 也确实有很多并发操作。那为什么在这个测试里没有看到不确定性呢？

为了理解 LLM 推理不确定性的真正原因，我们必须深入研究。

不幸的是，即使定义 LLM 推理的确定性意味着什么也很困难。一些令人困惑的事情是，以下所有陈述同时都是正确的：GPU 上的一些内核是不确定性的。

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

此外，LLM 推理引擎（如 vLLM）的前向计算也可以声称是确定性的。

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

在这篇文章中，我们将解释为什么“并发 + 浮点数”假设未能达到目标，揭示导致 LLM 推理不确定性背后的真正原因，并解释如何克服不确定性并在 LLM 推理中获得真正可重现的结果。

最初的“原罪”：浮点数的非结合性

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

罪魁祸首就是浮点数的非结合性。也就是说，对于浮点数：

(a + b) + c ≠ a + (b + c)

具有讽刺意味的是，正是这种对结合律的破坏使得浮点数变得有用。

浮点数之所以有用，是因为它们允许“动态”的精度级别。

为了便于解释，我们将使用十进制（而非二进制），其中浮点数采用 尾数 * 10^指数 的格式。我们还将使用 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 位精度。因此，浮点加法将丢弃最后两位，得到值 1.25 * 10^3（或 1250）。

Figure 1：我们需要 3 位精度来表示 1230，也需要 3 位精度来表示 23.4。然而，将这两个数相加会得到一个需要 5 位精度才能表示的数（1253.4）。我们的浮点格式必须丢弃末尾的 34。从某种意义上说，我们在相加之前，实际上已经将原始的 23.4 四舍五入到 20.0 了。

然而，在这一点上，我们已经丢失了信息。请注意，每当我们相加两个具有不同“尺度”（即不同指数）的浮点数时，都可能发生这种情况。而相加具有不同指数的浮点数是经常发生的。事实上，如果能保证从不需要不同的指数，我们就可以直接使用整数了！

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

尽管这是导致输出不一致的根本原因，但它并没有直接回答不确定性从何而来。它没有帮助我们理解为什么浮点值会以不同的顺序相加、何时发生这种情况以及如何避免。

答案在于内核（kernels）是如何实现的。

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

既然运算顺序会影响结果，那么为什么 GPU 内核在执行加法时，顺序不是固定的呢？

这就回到我们之前提到的一个流行假设：GPU 上的多个线程（并发执行）完成的顺序不确定，如果加法（累加）操作的最终结果依赖于这些不确定的完成顺序（例如，如果多个线程同时尝试向同一个内存位置累加，需要用到“原子加法”），那么最终的累加顺序就会不确定，从而导致结果不确定。

令人困惑的是，尽管这（并发与浮点非关联性）可能会导致内核出现不确定性，但在 LLM 推理的不确定性中，并发（以及原子加法 atomic adds）实际上完全没有参与其中！ 要解释真正的罪魁祸首是什么，我们首先需要理解为什么现代 GPU 内核很少需要使用原子加法。

*atomic add = 一种让很多线程同时往同一个变量里“加数”而不丢结果的机制，但加法顺序不固定，所以浮点数计算可能有细微差异。

什么时候需要用到原子加法？

通常情况下，GPU 会在许多“核心”（即 SM，Streaming Multiprocessors）上并发地启动一个程序。由于这些核心之间没有天然的同步机制，如果它们需要彼此通信，就会产生挑战。比如说，如果所有核心都必须把结果累加到同一个元素上，就可以使用“atomic add”（有时也叫“fetch-and-add”）。

atomic add 是“不确定性的” —— 结果的累加顺序完全取决于哪个核心先完成。

具体来说，假设你要用 100 个核心对一个 100 元素的向量做归约（reduction）（例如 torch.sum()）。虽然你可以并行加载这 100 个元素，但最终必须把它们归约成一个结果。一种实现方式就是使用某种 atomic add 原语，在这种情况下，硬件能保证所有加法都会被处理，但不能保证执行的顺序。

Figure2：atomic add（原子加法） 确保了每个核心的贡献都会体现在最终的总和中。然而，它并不保证这些贡献是按照什么顺序被加进去的。顺序完全取决于哪个核心先完成，这是一个不确定性特性。因此，执行同一个并行程序多次，可能会产生不确定性输出。

这通常就是人们所说的"不确定性"——用完全相同的输入执行同一个内核两次，却得到截然不同的结果。

这种现象被称为运行间不确定性（run-to-run nondeterminism）：你用相同的依赖环境运行同一个 Python 脚本两次，结果却不一样。

虽然并发的 atomic add 操作确实会让内核变得不确定，但对于绝大多数内核来说，atomic add 其实并非必需品。 事实上，在 LLM 的典型前向传播过程中，通常连一个 atomic add 操作都不会出现。

这个现象可能会让人感到惊讶——毕竟在对归约（reduction）操作进行并行化时，atomic add 确实能带来性能收益。 但 atomic add 最终之所以并非必需，主要有两个原因：

首先，在“batch（批次）”维度上通常已经有足够的并行性，因此我们并不需要在归约维度上再做并行化。举个例子，假设我们面对的不是对单个 100 维向量做归约，而是需要同时对 500 个向量做归约。在这种情况下，我们完全可以让每个核心负责处理一个完整的向量，让不同核心操作不同的向量，这样就自然地避免了竞争条件。

其次，随着时间的推移，大多数神经网络库都逐渐采用了多种巧妙的策略，在不牺牲性能的前提下实现了计算的确定性。比如，我们可以采用"分块归约（split reduction）"或"树状归约（tree reduction）"的方式：将一个 100 元素的归约任务拆分成 5 个独立的 20 元素归约（从而实现五路并行处理）。

接下来，为了将剩余的这 5 个结果合并起来，我们有两种选择：

执行一次单独的“清理归约（clean-up reduction）” —— 虽然这一步不再并行，但由于处理的元素数量很少，计算代价相当低廉；

或者采用信号量（semaphore）机制，它能够确保每个并发的线程块按照确定的顺序进行累加操作。

这样一来，我们既保持了高效的并行计算，又避免了不确定性带来的困扰。

由于以上两个因素，对于绝大多数神经网络运算来说，避免使用 atomic add 带来的性能损失几乎可以忽略不计。

当然，仍有一些常见操作在避免 atomic add 时会造成显著的性能损失，比如 PyTorch 中的 scatter_add 操作（a[b] += c）。不过在 LLM 领域，唯一常用到这类操作的场景就是 FlashAttention 的反向传播了。

这里有个有趣的事实：你知道吗？被广泛使用的 Triton 实现版本，在算法层面其实与 Tri Dao 的 FlashAttention-2 原版论文并不相同。

标准的 Triton 实现会在反向传播中进行额外的重计算，以此来规避 atomic add 的使用，但代价是增加了 40% 的 FLOPs！这是一个典型的以计算换确定性的权衡。

然而，LLM 的前向传播并不涉及需要 atomic add 的操作。因此，LLM 的前向传播实际上具备了运行间确定性（run-to-run deterministic）的特性。

从推理服务器的角度来看，这意味着什么呢？它是完全确定性的——只要用户的请求完全相同，服务器总是会返回相同的输出结果。

正如维基百科所定义的："一个确定性算法，是指在给定某个输入时，总是会产生相同输出的算法。"而 LLM 的前向推理恰好满足了这个定义。

在这里，给定完全相同的输入（即推理服务器正在处理的完全相同的请求），前向传播总是会产生完全相同的输出。

然而，仅仅前向传播本身具有"确定性"，并不足以保证包含它的整个系统也是确定性的。

举个例子，如果我们的请求输出会受到并行用户请求的影响（比如在 batch normalization 的情况下），那情况就大不相同了。

由于每个独立请求都无法预知其他并行请求的内容，从单个请求的视角来看，整个 LLM 推理系统仍然表现出不确定性！

事实上，我们的请求输出确实会受到并行用户请求的影响。

这种现象并非源于批次间的信息泄露，而是因为我们的前向传播缺乏批次不变性（batch invariance）——换句话说，请求的输出结果会依赖于前向传播时的具体批次大小。

批次不变性与确定性

为了说明批次不变性的问题，我们将系统简化，专门考察矩阵乘法（matmul）操作。

在讨论中，我们可以假设所有的矩阵乘法实现都具有"运行间确定性（run-to-run deterministic）"——尽管这个假设在实际情况中并非完全准确，但大多数常见的矩阵乘法实现确实表现出这种特性。

然而，一个令人意外的现象是：这些实现往往不具备"批次不变性（batch-invariant）"。也就是说，当批次大小发生变化时，批次中每个元素的计算结果可能会随之改变。

从数学的直觉来看，这确实是一个颇为反常的现象。按理说，矩阵乘法应该对批次中的每个元素"独立"进行计算——无论批次中包含其他什么元素，也无论批次的总体大小如何，都不应该影响到批次中任意特定元素的最终结果。

但令人困惑的是，实际的观察结果却与这种数学直觉相悖。

请注意，这里指的是"运行间确定性（run-to-run deterministic）"。如果你在同一环境中多次运行这个脚本，它会确定性地返回完全相同的结果。

不过，这种确定性是有条件的 —— 它并非"跨硬件/软件版本不变"。当你的GPU型号或PyTorch版本发生变化时，可能会得到不同的数值结果。但关键在于：在同一个运行环境下，结果始终是可重现的。

然而，问题出现在更复杂的场景中。当这种缺乏批次不变性的内核被集成到大型推理系统时，整个系统的行为就变得不可预测了。

想象一下这样的情况：当你向推理服务端点发送请求时，服务器当前的负载情况对你来说是"黑盒"的 —— 你无法预知也无法控制。而这个负载直接决定了内核运行时的批次大小，进而影响每个看似独立的请求的最终输出结果！这就是为什么同样的输入在不同时刻可能产生不同输出的根本原因。

Figure4：尽管从整体上可以说推理服务器是“确定性的”，但对单个用户来说情况不同。从单个用户的角度来看，其他并发用户并不是系统的“输入”，而是系统的一个不确定性属性。这使得从每个用户的角度来看，LLM 推理是“不确定性的”。

当某种内核缺乏批次大小不变性时，一旦遇到不确定的外部因素（比如服务器负载的波动），整个系统就会变得不确定起来。

实际上，几乎所有 LLM 推理端点的不确定性都可以追溯到同一个根源：服务器负载的动态变化会导致批次大小的不可预测变动！

这种不确定性并非 GPU 所独有——无论是基于 CPU 还是 TPU 的 LLM 推理端点，都会面临同样的挑战。

那么，如果我们希望在推理服务器中消除这种不确定性，就必须在内核层面实现批次不变性。要理解具体该如何实现，我们首先需要弄清楚：为什么内核最初就不具备批次不变性呢？

我们该如何让内核实现批次不变性？

为了让 Transformer 的实现具有批次不变性，我们需要确保其中的每一个内核都满足这个要求。

好消息是，我们可以安全地假设所有逐点操作（pointwise operation）都是批次不变的。

虽然这在 PyTorch 的所有内核中确实成立，但从理论上讲，这并非绝对保证。举个例子，在某些 CPU 内核实现中，可能会对数组的部分区域使用向量化指令（vectorized intrinsics），而对其他区域使用标量指令，这两种指令的数值结果未必总能做到逐位（bitwise）完全一致。

因此，我们真正需要关注的是三类涉及归约（reduction）的操作：RMSNorm、矩阵乘法（matrix multiplication）和注意力机制（attention）。

虽然与并行性相关的归约问题不在本次讨论范围之内，但相同的处理原则同样适用。

顺便分享一个实用的小知识：在 Blackwell 架构以及搭配 CUDA 12.8+ 的 Hopper 上，NVLink-Sharp 的交换机内归约（in-switch reductions）是确定性的。和很多技术细节一样，这个信息藏在 NCCL 的 GitHub issue 讨论中。

巧合的是，这三类操作的实现难度也恰好呈递增趋势。想要在保持合理性能的前提下实现批次不变性，每一个操作都需要我们投入额外的思考和设计。

那么，让我们先从 RMSNorm 开始讲起。

批次不变的 RMSNorm

Figure 5：理想情况下，我们希望在并行化策略中避免核心间的通信。实现这一目标的一种方法是将一个批次元素分配给一个核心，从而保证每次归约都完全在一个核心内完成。这就是所谓的“数据并行”策略，因为我们只是沿着一个不需要通信的维度进行并行化。在这个例子中，我们有四行和四个核心，充分利用了我们的核心。

对批量不变性的要求是：无论内核的批量大小是多少，每个元素的归约顺序（reduction order）都必须保持一致。

需要注意的是，这并不意味着我们必须始终使用完全相同的归约策略。举个例子，当要归约的元素数量发生变化时，即使归约策略相应调整，我们依然可以维持批量不变性。Quack 博客文章中提供了一些精彩的示例，清晰地展示了各种归约策略的层级关系 —— 从线程归约、warp 归约到 block 归约、cluster 归约。

换句话说，只有当批量大小的变化直接影响到归约策略的选择时，批量不变性才会被打破。

现在让我们来看看 RMSNorm 的标准并行策略。

众所周知，并行算法的性能优势主要来源于最大化减少核心之间的通信开销。为了便于理解，这里我们可以将"核心（cores）"简单理解为 SM（Streaming Multiprocessors，流式多处理器）。在实际应用中，有一个关键特性需要把握：我们启动的线程块（threadblocks）数量通常要超过可用的 SM 数量。基于这个前提，一个自然的策略就是为每个批量元素分配一个独立的核心，正如上图所展示的那样。

这种分配策略的美妙之处在于：增加批量大小并不会改变我们的归约策略。如果批量大小为 200 时已经能为内核提供充足的并行性，那么当批量大小增加到 2000 时，这种并行性只会变得更加充沛。

Figure 6：数据并行 RMSNorm（适用于更大的批次）将数据并行策略扩展到更大的批次相当直接 —— 与其让每个核心只处理一行，不如让每个核心依次处理不同的行。这样可以保持批次不变性，因为每个批次元素的归约策略仍然是相同的。

另一方面，减小批次大小可能会带来挑战。

因为我们为每个批次元素分配一个核心，当批次大小逐渐减小时，势必会遇到核心数量超过批次元素数量的情况，这时部分核心就会处于空闲状态。

面对这种情况，经验丰富的内核工程师通常会采用前一节介绍的某种解决方案（比如原子加法或分裂归约）来充分利用这些空闲核心，以保持良好的并行效率。然而，这样做的代价是改变了原有的归约策略，使得内核无法继续维持批次不变性这一重要特性。

Figure 7：分裂归约 RMSNorm。如果我们的推理批次较小，那么数据并行策略可能无法提供足够的并行度来充分利用所有核心。在这种情况下，一个更高效的做法是将单个归约任务分解到多个核心上并行处理，从而让 GPU 资源得到更好的利用。不过，这种方法会带来一个权衡：我们可能会失去批次不变性。这是因为当多个核心并行处理同一个归约任务时，元素的处理顺序可能不再保持一致，进而影响最终结果的确定性。

最简单的解决办法就是干脆完全忽略这些情况。

这其实也算合理 —— 毕竟小批次本身执行速度就很快，即便性能有所下降，通常也不至于造成严重问题。

如果确实需要优化这类场景，一个可行的思路是始终采用某种归约策略，确保即便在很小的批次下也能提供充足的并行度。虽然这种策略在大批次时会产生一定程度的过度并行，但它能让我们在各种批次规模下都获得相对稳定（尽管不是最优）的性能表现。

批次不变的矩阵乘法

Figure 8：数据并行矩阵乘法：与 RMSNorm 类似，矩阵乘法的标准并行策略采用"数据并行"的思路，即将整个归约过程集中在单个核心内完成。最直接的实现方式是：将输出张量划分为二维小块（2D tiles），每个小块分配给不同的核心处理。各核心负责计算所分配小块内的点积运算，归约操作仍在核心内部完成。然而，矩阵乘法相比 RMSNorm 面临着更多约束条件——算术强度（arithmetic intensity）的要求以及对 Tensor Cores 的高效利用等因素，都对内核设计提出了更严格的要求。正是这些约束，使得高效的矩阵乘法内核必须采用二维小块的切分策略，而非按单个输出元素进行切分的简单方式。

从本质上看，你可以把矩阵乘法理解为：先执行逐点运算 (pointwise operation)，再执行一次归约 (reduction)。

基于这个理解，当我们通过把输出切分成小块 (tiles) 来并行化矩阵乘法时，实际上得到了类似"数据并行"的内核策略——每次归约运算仍然保持在单个核心内部完成，这样就避免了跨核心的通信开销。

然而，就像我们在 RMSNorm 中遇到的问题一样，矩阵乘法的"批次"维度（M 和 N）有时会显得太小，迫使我们不得不沿着归约维度 (K) 进行切分。虽然矩阵乘法拥有两个"批次"维度的优势，但要想充分发挥 Tensor Cores 的威力，我们仍然需要在每个核心上分配足够多的"工作量"。

举个具体例子：假设你有一个形状为 [1024, K] × [K, 1024] 的矩阵乘法，如果采用标准的二维 tile 大小 [128, 128]，那么数据并行策略最多只能将这次计算分配给 64 个核心。对于现代 GPU 来说，这样的并行度显然还不足以让硬件"吃饱"。

在矩阵乘法领域，沿着归约维度进行切分的策略有一个专门的名字：Split-K Matmul。不过需要注意的是，和 RMSNorm 的情况类似，采用 Split-K 策略同样会破坏批次不变性 (batch invariance)——这是我们在选择并行策略时需要权衡的一个重要考量。

Figure 9：Split-K 矩阵乘法。如果我们的批次维度非常小，就可能没有足够的并行度，这时就需要使用 Split-K 矩阵乘法。在这个例子中，我们把每一次归约拆分给两个核心来执行，这两个核心会各自累加，最后再把结果合并。这样一来，虽然每次归约被分到两个核心上，但我们仍然能够利用八个核心。

矩阵乘法还有一个额外的复杂点 —— Tensor Core 指令。

在归约运算里，我们可以一次只处理一行；但在高效的矩阵乘法内核中，必须一次处理整个“tile”（小矩阵块）。

每条 Tensor Core 指令（比如

wgmma.mma_async.sync.aligned.m64n128k16）在内部可能会有不同的归约顺序。选择不同的 Tensor Core 指令的一个原因，可能就是 批次大小非常小。

举个例子：如果我们使用的 Tensor Core PTX 指令需要处理长度为 256 的 tile，但实际 batch size 只有 32，那几乎 90% 的算力都被浪费掉了！

当 batch size = 1 时，最快的内核通常甚至完全不会使用 Tensor Cores。

Figure10：填充 (Padded) Tensor Core 指令。如果批次大小太小，就可能出现一种情况：我们甚至无法在输出中放下一个完整的二维 tile。在这种情况下，最有效的做法是切换到更小的 Tensor Cores 指令，或者干脆完全不用 Tensor Cores！然而，这两种选择都会使我们的内核无法保持批次不变性 (batch-invariance)。

因此，确保矩阵乘法批次不变性最简单的方法，就是只编译一种内核配置，并把它应用到所有形状上。

虽然这样会损失一部分性能，但在大语言模型（LLM）的推理过程中，这通常并不是灾难性的。尤其是，Split-K 最需要的情况是 M 和 N 都很小，而幸运的是，在我们的场景中，N（也就是模型维度）通常都相当大！

Figure 11：尽管实现了批次不变性，我们相比 cuBLAS 只损失了大约 20% 的性能。需要注意的是，这里用的也不是经过优化的 Triton 内核（例如没有使用 Thinking MachinesA）。不过，一些性能表现的模式还是能说明问题，揭示出批次不变性要求导致性能损失的地方。首先，请注意在 非常小的批次大小 下，由于指令过大以及并行度不足，我们损失了相当多的性能。其次，当批次大小逐渐增加时，会出现一种 “拼图”式的性能波动模式，这是由量化效应（tile 和 wave 两方面）造成的，而这些效应通常可以通过改变 tile 大小来缓解。

批次不变的注意力机制

Figure 12：Flash Attention 2 策略。我们在 Q 维度 上进行并行，同时在 K/V 维度 上做归约。这意味着整个归约过程都能保持在单个核心内完成，使其成为另一种数据并行策略。

在矩阵乘法实现了批次不变性之后，注意力机制又引入了两个额外的复杂点 —— 这也合乎逻辑，因为它本身包含了两次矩阵乘法。

1. 与 RMSNorm 和矩阵乘法只在 特征维度 上做归约不同，注意力机制需要同时在 特征维度 和 序列维度 上做归约。

2. 因此，注意力机制必须应对多种推理优化方式，这些方式会影响序列的处理方式（例如 分块预填充 chunked prefill、前缀缓存 prefix caching 等）。

因此，要在大语言模型 (LLM) 的推理中实现确定性，我们的数值必须对两方面保持不变：既要与一次同时处理多少请求无关，也要与推理引擎如何切分每个请求无关。

接下来，我们先来看注意力机制的标准并行化策略，这个策略最早由 Flash Attention 2 提出。与 RMSNorm 和矩阵乘法类似，默认策略是一种“数据并行”策略。由于我们在 Key/Value 张量 上进行归约，所以数据并行策略只能在 Query 张量 上实现并行。

举个例子：根据推理引擎的不同选择，一个序列可能被分成多个部分来处理（比如分块预填充），也可能一次性处理全部（如果预填充没有被拆分）。为了实现“批次不变性”，关键在于：某个 token 的归约顺序不能依赖于该序列中同时处理了多少其他 token。

如果你把 KV 缓存里的 K/V 值与当前处理的 token 的 K/V 值分开做归约（就像 vLLM 的 Triton 注意力内核那样），那就无法做到批次不变性。

举例来说：当处理序列中的第 1000 个 Query token 时，归约顺序必须完全一致——不论 KV 缓存中有 0 个 token（预填充阶段）还是有 999 个 token（解码阶段）。

Figure 13：带 KV 缓存的 Flash Attention将 KV 缓存与当前 KV 值分开处理之所以会破坏批次不变性，这一点有些微妙，主要与 “边界条件 (boundary conditions)” 有关。

举个例子：假设块大小是 32，而我们当前在 KV 缓存中有 80 个元素。接着我们又计算了额外 48 个不在缓存中的元素。

在这种情况下：

计算 P cache 需要 3 个块（两个完整块 + 一个带掩码的块）。

计算 P 需要 2 个块（一个完整块 + 一个带掩码的块）。

这样一来，总共需要 5 个块来完成归约，但实际上我们只有 128 个元素，即 4 个块。这就必然会改变归约的顺序。

对比：如果我们一开始没有 KV 缓存，而是一次性处理 128 个元素，那么只需要 4 个块来完成归约。因此，为了保证注意力机制的 批次不变性 (batch invariance)，我们必须确保这两种情况下的数值结果完全一致。

为了解决这个问题，我们可以在进入注意力内核之前，先更新 KV 缓存 和 页表 (page table)，以确保无论当前处理多少个 token，我们的 Key 和 Value 都始终以一致的方式排列。

结合这一点（以及上一节提到的所有细节，比如保持一致的 tile 大小），我们就能够实现一个 批次不变的注意力机制 (batch-invariant attention)！

然而，这里存在一个严重的问题。不同于矩阵乘法，在大语言模型 (LLM) 推理中遇到的注意力形状往往确实需要 分裂归约内核 (split-reduction kernel)，通常称为 Split-KV 或 Flash Decoding。

原因是：如果我们不在归约维度上并行，就只能在 批次维度 (batch dimension)、头部维度 (head dimension) 和 Query 长度维度 (query length dimension) 上进行并行。

但在注意力的 解码阶段 (decode stage)，Query 长度非常小，所以除非批次很大，否则我们经常无法让 GPU 得到充分利用。

遗憾的是，这种情况不像 RMSNorm 和矩阵乘法那样可以轻松忽略。举个例子：如果 KV 缓存非常长，即便我们只在处理一个请求，注意力内核也可能会耗费非常长的时间。

Figure 14：固定的 # Split-KV 策略（也叫 Flash Decode）。如果我们的 query 长度变得非常小（就像在解码阶段那样），我们可能会遇到这样一种情况：内核中几乎没有多少并行度可利用。在这些情况下，我们需要再次沿着归约维度进行拆分 —— 这一次是 KV 维度。

典型的做法是先确定我们需要多少并行度，然后将 KV 维度均匀地划分。比如，如果 KV 长度是 1000，而我们需要 4 个拆分，那么每个核心就会处理 250 个元素。

但不幸的是，这种方法也会破坏 批量不变性（batch invariance），因为我们精确的归约策略取决于当前请求中要处理多少个序列的 query token。

此外，注意力中常用的 拆分归约（split-reduction）策略 也会给批量不变性带来挑战。举例来说，FlashInfer 的“平衡调度算法（balanced scheduling algorithm）”会选择一种最大化的拆分规模，以便仍能让 GPU 的所有核心饱和，因此这种归约策略并不是“批量不变的（batch-invariant）”。不过，与 RMSNorm/矩阵乘法不同，仅仅选择一个固定的拆分数量，而不管批量大小是多少，是不够的。

为了实现批量不变性，我们必须采用“固定拆分大小（fixed split-size）”的策略。换句话说，我们不是固定拆分的数量，而是固定每个拆分的大小，最终得到的拆分数量则会随之变化。这样一来，我们就能保证无论处理多少 token，始终执行相同的归约顺序。

Figure 15：固定大小的 Split-KV 策略这种策略与前一种策略的唯一不同在于，我们的拆分现在是“固定大小”的。比如，如果 KV 长度是 1000，那么我们不会再把它均分成四段、每段长度 250，而是拆分成三段固定长度 256 的部分，以及一段长度 232 的部分。

这样，我们就能够保持 批量不变性（batch invariance），因为我们的归约策略不再依赖于一次要处理多少个 query token！

实现

我们基于 vLLM，利用其 Flex Attention 后端以及 torch.Library，提供了一个确定性推理的演示。通过 torch.Library，我们能够以一种非侵入的方式替换掉大部分相关的 PyTorch 算子。你可以在 thinking-machines-lab/batch-invariant-ops 中找到“批量不变（batch-invariant）”内核的库，以及在 vLLM 中运行“确定性（deterministic）”模式的示例。

实验

补全（completions）的不确定性有多大？

我们使用 Qwen/Qwen3-235B-A22B-Instruct-2507，在温度设为 0 的情况下采样 1000 次补全，提示词为 “Tell me about Richard Feynman”（非思考模式），每次生成 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 中 FlexAttention 的集成尚未经过深度优化。尽管如此，我们可以看到性能并没有糟糕到不可接受的程度。

满血版 On-Policy 强化学习（RL）

正如研究者指出的那样，训练与推理阶段数值的不一致，实际上会把我们的 on-policy RL 变成 off-policy RL。

当然，如果连两次完全相同的推理请求都无法得到逐位（bitwise）一致的结果，那么训练和推理之间就更不可能实现逐位一致。而一旦实现了 确定性推理，我们也就能够修改训练栈，使得采样和训练之间的结果逐位一致，从而实现真正的 on-policy RL。

我们在 BigMath 上运行了一个 RLVR 设置实验，RL 策略从 Qwen 2.5-VL instruct 8B 初始化，最大采样长度为 4096。

如果我们在训练中不做 off-policy 修正（即：重要性加权，importance weighting），奖励会在训练过程中途崩溃；而在训练中加入 off-policy 修正项，训练则能顺利进行。但如果我们能够让采样器和训练器之间的结果逐位一致，那么我们就是完全的 on-policy（即 KL 散度为 0），同样也能顺利完成训练。

我们还可以绘制 采样器与训练器之间 logprob 的 KL 散度曲线，结果显示三种运行方式的行为有明显不同：

使用 重要性加权 时，KL 散度大约维持在 0.001，偶尔会有尖峰；

不使用 重要性加权 时，KL 散度最终会出现尖峰，并与奖励崩溃的时间点吻合；

而在运行 “真正的 On-Policy RL” 时，KL 散度始终保持在 0，表明训练策略与采样策略之间不存在任何差异。

Figure 16：请注意，在没有使用 重要性加权 的运行中，大约在 第 318 步时出现了显著的损失尖峰，同时伴随着 logprob 的 KL 散度的尖峰。与此同时，不论是使用 off-policy 修正，还是运行 “真正的 On-Policy”，都能让 RL 顺利地继续进行。图中那条代表 “True On-Policy” 的蓝色线并不是 bug —— 它只是一直平坦地保持在 0。

结论

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

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

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

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