成立7个月首发声！百亿美金独角兽万字雄文：攻克LLM推理非确定性难题

Thinking Machines Lab成立7个月，估值120亿美元，首次公开研究成果：LLM每次回答不一样的真凶——kernel缺乏批处理不变性。Lilian Weng更是爆猛料：首代旗舰叫 Connection Machine，还有更多在路上。

Thinking Machines Lab终于放大招了！

刚刚，联合创始人、OpenAI前副总裁Lilian Weng透露：

Thinking Machines的第一代旗舰产品名为「Connection Machine」（联结机）。

事情是这样的：今天，Thinking Machines开辟了研究博客专栏「Connectionism」（联结主义），发表了第一篇博客文章「Defeating Nondeterminism in LLM Inference」（击败LLM推理中的非确定性）。

Thinking Machines介绍说:

我们相信科学因分享而更美好。

Connectionism专栏将随着我们的研究变化：从内核数值计算到提示工程。在这里，我们分享我们的工作进展，并与研究社区频繁而开放地交流。

此外，还补充道，「Connectionism」这一名称可追溯至早期的AI时代——在1980年代，该术语指代专门研究神经网络及其与生物大脑相似性的子领域。

而Lilian Weng爆出了更大的料，专栏之所以叫这个名称，还有一个原因：第一代旗舰模型就叫Connection Machine，不光是这篇博客文章，而且还有更多好东西要来了！

莫非Thinking Machines马上要发布新模型了？

在期待新的LLM之前，我们先看看这次Thinking Machines到底有哪些绝活，他们到底关注哪些研究领域。

传送门：https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/

据博文的主要作者Horace He介绍，这次的博文主要关于他心中的重要话题——

LLM推理中的可复现浮点数（Reproducible floating point numerics in LLM inference）。

LLM推理中的非确定性难题

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

例如，你可能会观察到，多次向ChatGPT提出相同的问题会得到不同的结果。

这本身并不奇怪，因为从语言模型获得结果涉及「采样」：

将语言模型的输出转换为概率分布，概率性地选择一个token。

更令人惊讶的可能是，即使我们将temperature降至0（从而使采样在理论上是确定性的），LLM API在实践中仍然不是确定性的。

即使在自己的硬件上，使用像vLLM或SGLang这样的开源推理库运行推理，采样仍然不是确定性的。

但为什么LLM推理引擎不是确定性的呢？

一个常见的假设是，浮点数非结合性与并发执行的某种组合，导致了基于哪个并发核心先完成的非确定性。

这次研究则称之为LLM推理非确定性的「并发+浮点数」假说。

例如，华人研究员Jiayi Yuan、Hao Li、Xinheng Ding等最近上传了一篇arXiv预印本，其中写道：

GPU中的浮点运算表现出非结合性，意味着 (a+b)+c ≠ a+(b+c)，这是由于有限的精度和舍入误差。

此属性直接影响Transformer架构中注意力分数和logits的计算，其中跨多个线程的并行操作可能会根据执行顺序产生不同的结果。

传送门：https://arxiv.org/abs/2506.09501

虽然这个假说有些道理，但并未揭示全貌。

例如，即使在GPU上，对相同数据重复运行相同的矩阵乘法，也总会提供逐位相等的结果。

我们确实在使用浮点数，GPU确实有大量的并发计算。那为什么在这个测试中没有看到非确定性呢？⬇️

要理解LLM推理非确定性的元凶，我们必须更深入地探究。

不幸的是，即使是定义LLM推理的确定性也并非易事。

也许令人困惑的是，以下陈述竟能同时成立：

• GPU上的某些核（Kernel）是非确定性的。
• 然而，语言模型前向传播中使用的所有Kernel都是确定性的。
• 此外，LLM推理服务（如vLLM）的前向传播也可以声称是确定性的。
• 然而，从推理服务的用户的角度来看，结果都是非确定性的。

这次，Thinking Machines决定要揭示LLM推理非确定性背后的真正元凶，并阐述如何克服LLM推理中的非确定性，获得真正可复现的结果。

关键发现：

LLM前向传播不需要原子加法；其非确定性真正来源是「批次大小变化」而非「原子竞争」。

要想在推理服务中避免非确定性、为了使Transformer实现具有批处理不变性，我们必须在kernel中实现「批处理不变性」。

幸运的是，我们可以假设每个逐点（pointwise）操作都具有批处理不变性。因此，只需要担心涉及归约的3个操作——RMSNorm、矩阵乘法和注意力。

它们的实现难度也是递增的。每个操作都需要一些额外的考虑，才能以合理的性能实现批处理不变性。

批处理不变的RMSNorm： 数据并行RMSNorm

理想情况下，我们希望在并行化策略中避免核心之间的通信。

实现这一点的一种方法是为每个核心分配一个批处理元素，从而保证每个归约都完全在单个核心内完成。

这就是所谓的「数据并行」策略，因为我们只是沿着一个不需要通信的维度进行并行化。

批处理不变的矩阵乘法：数据并行Matmul

与RMSNorm类似，矩阵乘法的标准并行策略是一种「数据并行」策略，将整个归约保持在一个核心内。

最直接的思考方式是将输出张量分割成二维的分块（tiles），并将每个分块分配给不同的核心。然后，每个核心计算属于该分块的点积，再次在单个核心内执行整个归约。

与RMSNorm不同，围绕算术强度和利用张量核心（tensorcores）的额外约束，被迫分割二维分块而不是单个输出元素进行，以实现高效的矩阵乘法kernel。

解决的核心在于，你可以将矩阵乘法看作是一个逐点操作后跟一个归约。

确保矩阵乘法具有批处理不变性的最简单方法是，编译一个kernel配置，并将其用于所有形状。

虽然会损失一些性能，但这在大语言模型推理中通常不是灾难性的：

相比cuBLAS只损失了约20%的性能。

批处理不变的注意力机制

在为矩阵乘法获得批处理不变性之后，注意力机制引入了两个额外的难题——恰如其分，因为它包含两个矩阵乘法。

• 与RMSNorm和矩阵乘法仅在特征维度上进行归约不同，现在在特征维度和序列维度上进行归约。

• 由于上述原因，注意力机制必须处理各种影响序列处理方式的推理优化（分块预填充chunked prefill、前缀缓存prefix caching等）。

带KV缓存的FlashAttention会破坏批处理不变性，根因在把「缓存KV」与「当前KV」分开算

不同 KV 块数 → 不同掩码/完整块组合 → 不同规约路径。

只要在kernel启动前，统一更新KV-cache页表，保证任意时刻KV布局一致，就能解决这一问题。

大语言模型推理中看到的注意力形状通常确实需要一个分裂归约的kernel，通常称为Split-KV或FlashDecoding。

固定数量的Split-KV策略（即FlashDecode），因为精确的归约策略取决于给定请求中处理来自序列的查询token数量，这不幸地也破坏了批处理不变性。

如果我们的查询长度变得非常小（就像在解码期间那样），可能会陷入一种情况，即kernel中几乎没有并行性。在这些情况下，需要再次沿着归约维度——这次是KV维度——进行分割。分割KV维度的典型策略是计算出需要多少并行性，然后均匀地划分KV维度。例如，如果KV长度是1000，我们需要4个分割split，每个核心将处理250个元素。

此外，通常用于注意力的分裂归约策略也对批处理不变性构成了挑战。

为了实现批处理不变性，不再固定分割的数量，而是固定每个分割的大小，然后得到一个可变数量的分割。

通过这种方式，可以保证无论正在处理多少个token，我们总是执行相同的归约顺序。

这实现了批处理不变性，因为归约策略不再依赖于一次处理多少个查询token！

用「固定块大小」Split-KV，注意力也能像 RMSNorm/Matmul 一样实现批处理不变，确定性推理。

开实现与实验

通过利用vLLM的FlexAttention后端以及torch.Library，Thinking Machines提供了一个在vLLM上进行确定性推理的演示。

传送门：https://github.com/thinking-machines-lab/batch_invariant_ops

补全结果有多大的非确定性？

使用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」。

另一方面，当启用批处理不变kernel时，所有的1000个补全都是完全相同的。

性能

这次没有投入大量精力来优化批处理不变kernel的性能，但仍用实验测试了一下性能。

实验设置：一个带有一块GPU的API服务器，运行Qwen-3-8B，并请求1000个序列，输出长度在90到110之间。

大部分的性能下降来自于vLLM中的FlexAttention集成尚未经过大量优化。尽管如此，性能下降并非不可接受。

真正的同策略强化学习

正如研究人员指出的，训练和推理之间不同的数值计算，无形中将同策略强化学习（on-policy RL）变成了异策略强化学习（off-policy RL）。

传送门：https://fengyao.notion.site/off-policy-rl

如果两次相同的推理请求都无法做到逐位（bitwise）一致，那训练与推理在位级一致就更无从谈起。

确定性推理让我们可以同步改造训练栈，使采样与训练在数值上逐位一致，从而获得真正的同策略RL。

在Bigmath的RLVR设定下，研究人员做了实验：策略以Qwen 2.5-VL instruct 8B初始化，最大rollout长度 4096。

无异策略校正（不做重要性加权）时，训练中段奖励出现崩塌；

加入异策略校正（importance weighting）后，训练可平稳推进；

当让采样器与训练器逐位一致时，策略完全同源（KL=0），训练同样平稳。

同时，研究人员绘制了采样器与训练器对数概率（logprobs）之间的 KL 散度：三种设定差异明显——

• 带重要性加权(下图中的红线)时，KL约在 0.001附近，偶有尖峰；
• 不加权(下图中的绿线)时，KL 在奖励崩溃前后出现明显飙升；
• 真正同策略(下图中的蓝线)时，KL始终为0，表明训练策略与采样策略无分歧。

需要强调的是：未做重要性加权的那次运行在Step 318左右出现显著的损失峰值，同时KL散度同步陡升；而做了异策略校正或实现「真正同策略」的两种设置，RL都能持续、平滑地优化。

万事怕认真

现代软件系统层层抽象，机器学习中的非确定性与微小数值差异，往往让人想「睁一只眼闭一只眼」：

反正系统本就「概率化」，多一点不确定也无妨？

单元测试里把atol/rtol往上调一调、把训练与采样间的logprob差异当成「假阳性」，似乎也能过关。

请拒绝这种「算了吧」的心态。只要多做一点功课，我们完全可以定位并修复这些非确定性根源！

Thinking Machines希望本文能为社区提供一套解决推理端非确定性的清晰思路，也能激励更多人真正吃透自己的系统。

参考资料：

https://x.com/lilianweng/status/1965828743152509198

https://x.com/cHHillee/status/1965828670167331010

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

文章来自于“新智元”，作者“KingHZ”。