Defeating Nondeterminism in LLM Inference

文章标题：击败LLM推理中的不确定性
作者/机构：Horace He and Thinking Machines Lab

A1 主要贡献

本文探讨并解决了大型语言模型（LLM）推理中普遍存在的不确定性问题。即使在理论上应为确定性的贪婪采样（temperature=0）模式下，LLM的输出在实践中仍然不一致。

核心问题：LLM推理引擎为何不是确定性的？流行的“并发+浮点数”假说认为，GPU中浮点数运算的非结合性与并发线程的执行顺序不确定性相结合，导致了结果的差异。然而，本文指出这个假说并不完全准确，因为它无法解释为何某些GPU操作（如矩阵乘法）在重复执行时能得到逐比特相同的结果。

研究目标：本文旨在揭示LLM推理不确定性的真正根源，并提出一种能够实现完全可复现结果的解决方案。

创新与贡献：
1. 揭示了不确定性的真正原因：文章论证了LLM推理不确定性的主要原因并非来自单个内核（Kernel）的“运行时不确定性”（run-to-run nondeterminism），因为LLM前向传播中使用的内核本身是确定性的。真正的罪魁祸首是这些内核缺乏“批处理不变性”（batch invariance）。这意味着，即使对于批处理中的同一个请求，其计算结果也会因为整个批处理的大小（batch size）不同而发生改变。由于推理服务器的负载是动态变化的，导致批处理大小不确定，进而使得最终输出不确定。
2. 提出了实现确定性推理的系统性方法：为了解决这个问题，文章提出必须使推理过程中的所有关键操作都具备批处理不变性。作者详细分析了实现以下三个关键操作批处理不变性的策略和挑战：
* RMSNorm：通过采用固定的、不随批处理大小变化的归约策略。
* 矩阵乘法（Matrix Multiplication）：通过编译和使用单一的、固定的内核配置（包括固定的分块大小和张量核心指令），避免使用如Split-K这类会根据输入形状改变计算方式的优化。
* 注意力机制（Attention）：通过在注意力计算前统一更新KV缓存，并采用“固定分块大小”（fixed split-size）而非“固定分块数量”的归约策略（如Split-KV），来保证无论序列如何被切分处理，其归约顺序都保持一致。
3. 提供了开源实现和实验验证：作者基于vLLM框架，利用FlexAttention后端和torch.Library，实现了一个支持确定性推理的原型。实验证明，该方法能让LLM在多次重复请求下生成完全相同的文本。虽然存在一定的性能开销，但仍在可用范围内。此外，文章还展示了确定性推理在实现“真正的在线策略强化学习”（True on-policy RL）中的重要价值，避免了因训练和采样阶段的数值差异导致的策略偏离。

A3 背景知识与问题分析

根本原因：浮点数的非结合性

浮点数非结合性的本质。机器学习模型通常被视为遵循交换律或结合律等数学规则的函数，但浮点数运算打破了这一常规。问题的根源在于浮点数的非结合性，即 $(a+b)+c \neq a+(b+c)$。例如，(0.1 + 1e20) - 1e20 的结果是 0，而 0.1 + (1e20 - 1e20) 的结果是 0.1。这种特性使得浮点数非常有用，因为它允许动态调整精度。浮点数以 $mantissa \times 10^{exponent}$ 的形式表示，可以同时表示极大和极小的数值，保持恒定的“有效数字”位数。

浮点数非结合性示例
浮点数非结合性示例

信息丢失的发生机制。当两个指数不同的浮点数相加时，问题就出现了。例如，1230 和 23.4 相加，精确结果是 1253.4。但如果只能保持3位精度，浮点加法会舍弃最后两位，得到 1.25 \times 10^3（即1250），这导致了信息丢失。这种情况在每次对不同“尺度”（即不同指数）的浮点数进行加法时都可能发生。因此，每次以不同顺序对浮点数求和，都可能得到完全不同的结果。

求和顺序导致结果多样性。为了证明这一点，一个极端的例子是，对一个包含正负成对数值的数组进行求和，仅仅因为求和顺序的不同，就可能产生102种不同的结果。这揭示了数值差异的根本原因，但并未直接解释不确定性的来源，即为何浮点数会以不同的顺序相加。答案在于内核（kernel）的具体实现方式。

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推理的不确定性问题中，并发和原子加法（atomic adds）并非根本原因。

原子加法的必要性。当多个GPU核心需要对同一个元素进行累加时，可以使用“原子加法”指令。该指令能保证所有加法操作都被处理，但不保证处理顺序，因此其结果依赖于核心的完成顺序，从而产生“运行时不确定性”（run-to-run nondeterminism）。这意味着，用完全相同的输入两次执行同一个内核，会得到不同的结果。

LLM前向传播中原子加法的缺失。尽管原子加法会导致内核不确定，但绝大多数内核并不需要它，LLM的典型前向传播中通常不存在任何原子加法。这主要有两个原因：
1. 批处理维度的并行性：通常在“批处理”维度上有足够的并行性，因此无需在归约（reduction）维度上进行并行化。例如，若要并行处理500个向量的归约，可以将每个向量分配给一个核心独立处理。
2. 确定性归约策略的采用：随着时间推移，大多数神经网络库采用了各种策略来实现确定性而不牺牲性能。例如，可以执行“分裂”或“树状”归约，将一个大的归约任务分解成多个小的并行任务，最后用一个小的、非并行的“清理”归约或使用信号量（semaphore）来确保以确定性顺序合并结果。

结论：LLM前向传播是运行时确定性的。由于上述两个因素，避免使用原子加法对绝大多数神经网络操作的性能影响微乎其微。虽然某些操作（如PyTorch中的scatter_add）和FlashAttention的反向传播确实会因避免原子操作而产生显著性能损失，但这些在LLM的前向传播中并不涉及。因此，LLM的前向传播本身是“运行时确定性”的。给定完全相同的输入，前向传播总会产生完全相同的输出。

A2 方法细节

批处理不变性与“确定性”

从运行时确定性到系统不确定性。尽管LLM前向传播本身是“运行时确定性”的，但这并不足以保证整个推理系统是确定性的。如果一个请求的输出依赖于与其并行的其他用户请求（例如，batch-norm），那么从单个请求的角度看，整个LLM推理过程就是不确定性的。事实证明，请求的输出确实依赖于并行的用户请求，但并非因为批次间的信息泄露，而是因为前向传播缺乏“批处理不变性”（batch invariance），导致请求的输出依赖于前向传播的批处理大小。

批处理不变性的缺失。为了解释批处理不变性，我们可以仅关注矩阵乘法。我们可以假设所有矩阵乘法的实现都是“运行时确定性”的，但它们并非“批处理不变”的。换言之，当批处理大小改变时，批处理中的每个元素都可能得到不同的计算结果。从数学角度看，这是一个不寻常的属性，因为矩阵乘法在批处理的每个元素上应该是独立的。

实证演示。以下代码凭经验证明了这一点。当分别计算一个批次中的第一个元素的矩阵乘法，与先计算整个批次的矩阵乘法再取出第一个元素的结果相比，两者存在显著差异。尽管这个脚本每次运行时都会确定性地返回相同的结果（即“运行时确定性”），但结果本身是依赖于计算时的“批处理大小”的。

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)   
# 通过取批次的第一个元素进行矩阵向量乘法
out1 = torch.mm(a[:1], b)   
# 进行矩阵乘法然后取批次的第一个元素
out2 = torch.mm(a, b)[:1]   
print((out1 - out2).abs().max()) # tensor(1669.2500, device='cuda:0')

系统不确定性的根源。当一个不具备批处理不变性的内核被用于一个更大的推理系统时，该系统就会变得不确定。用户向推理端点发出查询时，服务器的负载量是“不确定”的。负载决定了内核运行时所采用的批处理大小，从而改变了每个独立请求的最终结果。因此，将内核不具备不变性的属性（如批处理大小）与该属性的不确定性（如服务器负载）相结合，便得到了一个不确定的系统。几乎所有LLM推理端点不确定性的主要原因，就是负载（从而导致批处理大小）的不确定性变化。这种不确定性与硬件无关，在CPU或TPU上同样存在。要解决这个问题，就必须在内核中实现批处理不变性。

如何使内核具备批处理不变性？

核心任务。要使Transformer实现批处理不变性，必须使其每个内核都具备该特性。幸运的是，我们可以假设所有逐点（pointwise）操作都天然具备批处理不变性。因此，我们只需关注涉及归约（reduction）的三个操作：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

批处理不变性的要求。批处理不变性的要求是，对于每个元素的归约顺序必须固定，无论内核的批处理大小如何。这并不意味着必须始终使用同一种归约策略。例如，即使归约的元素数量改变，只要归约策略保持不变，仍然可以实现批处理不变性。因此，只有当批处理大小影响到归约策略时，才会破坏批处理不变性。

标准并行策略及其问题。RMSNorm的标准并行策略是为每个批处理元素分配一个核心（SM）。当增加批处理大小时，这种策略不会影响归约方式。然而，当减小批处理大小时，可能会出现核心数多于批处理元素数的情况，导致部分核心闲置。在这种情况下，内核工程师通常会采用原子加法或分裂归约等方案来维持高性能，但这会改变归约策略，从而破坏批处理不变性。

解决方案。最简单的解决方案是直接忽略小批处理大小下的性能优化，因为此时内核执行速度本就很快，性能下降可能不严重。如果必须优化，可以采用一种固定的归约策略，该策略即使在非常小的批处理下也能提供足够的并行度。虽然这种策略在较大批处理下会产生过多的并行性，导致性能非最优，但它能在整个尺寸范围内实现可接受的性能并保证批处理不变性。

批处理不变的矩阵乘法

矩阵乘法与归约。矩阵乘法本质上可以看作一个逐点操作后跟一个归约操作。如果我们将矩阵乘法通过对输出进行分块（tiling）来并行化，就得到了一种类似于RMSNorm的数据并行内核策略，这种策略将每个归约操作限制在单个核心内。

并行策略的挑战。与RMSNorm类似，当矩阵乘法的“批处理”维度（M和N）变得过小时，为了有效利用Tensor Cores，我们可能被迫沿归约维度（K）进行切分。例如，一个 [1024, K] x [K, 1024] 的矩阵乘法，若使用 [128, 128] 的标准2D分块，数据并行策略只能将其拆分到64个核心上，不足以饱和整个GPU。

Split-K策略与不变性破坏。在矩阵乘法中沿归约维度进行切分被称为Split-K Matmul。与RMSNorm中的情况一样，使用这种策略会破坏批处理不变性。另一个有趣的并行策略是Stream-K，它甚至比典型的矩阵乘法具有更少的不变性。大多数矩阵乘法库虽然不具备批处理不变性，但至少具备“批处理位置不变性”（即改变元素在批次中的位置不影响其数值）。然而，Stream-K的核心思想是为不同的输出块采用不同的K维度切分方式以实现更好的负载均衡，这使其连批处理位置不变性也无法保证。

Tensor Core指令的影响。矩阵乘法还存在一个额外的复杂性：Tensor Core指令。高效的矩阵乘法内核必须一次处理整个“分块”。不同的Tensor Core指令（如 wgmma.mma_async.sync.aligned.m64n128k16）内部可能有不同的归约顺序。当批处理大小非常小时，内核可能会选择不同的Tensor Core指令，甚至完全不使用Tensor Core，以避免计算资源的浪费，但这同样会破坏批处理不变性。

解决方案。确保矩阵乘法批处理不变性的最简单方法是，为所有形状编译并使用同一个内核配置。尽管会损失一些性能，但这在LLM推理中通常不是灾难性的。特别地，Split-K主要在M和N维度都很小时才需要，而在LLM推理中，N维度（即模型维度）通常足够大。

批处理不变的注意力机制

注意力机制的额外复杂性。在实现了矩阵乘法的批处理不变性后，注意力机制引入了两个额外的难题，因为它包含两次矩阵乘法：
1. 它不仅像RMSNorm和矩阵乘法那样在特征维度上进行归约，还需要在序列维度上进行归约。
2. 因此，注意力机制必须处理各种影响序列处理方式的推理优化（如分块预填充、前缀缓存等）。
要实现LLM推理的确定性，数值计算必须对一次处理多少请求以及每个请求如何被切分都保持不变。

标准并行策略与挑战。注意力机制的标准并行策略（由FlashAttention-2引入）是一种“数据并行”策略，在键/值（K/V）张量上进行归约，因此只能在查询（Q）张量上进行并行化。为了实现“批处理不变性”，对于一个给定的Token，其归约顺序不能依赖于其序列中同时被处理的其他Token的数量。如果像vLLM的Triton注意力内核那样，将KV缓存中的K/V值与当前正处理的Token的K/V值分开进行归约，就无法实现这一点。例如，在处理序列中的第1000个查询Token时，无论KV缓存中有0个Token（预填充阶段）还是999个Token（解码阶段），其归约顺序都必须完全相同。

初步解决方案。要解决这个问题，我们可以在注意力内核本身执行之前，先更新KV缓存和页表，从而确保无论正在处理多少Token，键和值总是以一致的方式布局。

Split-KV策略与最终解决方案。然而，一个重要问题依然存在。与矩阵乘法不同，LLM推理中的注意力形状通常确实需要分裂归约内核（常被称为Split-KV或FlashDecoding），尤其是在解码阶段，查询长度非常小，若不沿归约维度并行，将无法饱和GPU。不幸的是，常用的分裂归约策略也给批处理不变性带来了挑战。例如，FlashInfer的“平衡调度算法”会选择能够饱和GPU所有核心的最大分裂尺寸，这使得归约策略不具备批处理不变性。要实现批处理不变性，我们必须采用“固定分块大小”（fixed split-size）策略。即，我们固定每次归约的分块大小，而不是固定分裂的数量。这样，无论处理多少Token，我们总能执行相同的归约顺序。这需要对FlexAttention进行一些内部修改。

A7 补充细节（实现）

实现方式。我们在vLLM之上，通过利用其FlexAttention后端以及torch.Library，提供了一个确定性推理的演示。通过torch.Library，我们能够以非侵入性的方式替换掉大多数相关的PyTorch算子。我们提供了一个名为“batch-invariant-ops”的库，包含了这些具备“批处理不变性”的内核，以及一个在vLLM中以“确定性”模式运行的示例。你可以在thinking-machines-lab/batch-invariant-ops找到这个库。

A4 实验

实验环境

  • 模型：
    • Qwen/Qwen3-235B-A22B-Instruct-2507（用于完成生成实验）
    • Qwen-3-8B（用于性能测试）
    • Qwen 2.5-VL instruct 8B（用于RL实验）
  • 硬件配置：
    • 单GPU（具体型号未说明）
  • 软件配置：
    • vLLM框架
    • FlexAttention后端
    • torch.Library
    • 自定义的batch-invariant-ops库
  • 数据集/任务：
    • 完成生成：使用提示 "Tell me about Richard Feynman" 生成1000个Token的补全。
    • 性能测试：API服务器，处理1000个序列请求，输出长度在90到110之间。
    • 强化学习：在Bigmath数据集上进行RLVR（Reinforcement Learning from Video Representations）设置的实验，最大 rollout 长度为4096。

实验结果

补全的不确定性有多大？
* 实验内容：使用Qwen-235B模型，在temperature=0的设置下，对同一提示生成1000次，每次生成1000个Token。
* 实验结果：在标准模式下，令人惊讶地生成了80个独特的补全结果，最常见的结果也只出现了78次。所有补全在前102个Token上是相同的，但在第103个Token处开始出现分歧。例如，992个补全生成了“Queens, New York”，而另外8个生成了“New York City”。
* 结论：启用批处理不变内核后，1000次生成的所有补全结果都是完全相同的，符合确定性采样的数学预期。

性能
* 实验内容：设置一个单GPU的API服务器，运行Qwen-3-8B模型，处理1000个序列请求。
* 实验结果：与标准vLLM相比，启用确定性模式后，吞吐量有所下降，首次Token延迟和Token间延迟均有增加。作者指出，性能下降很大一部分原因是vLLM中FlexAttention的集成尚未经过深度优化。
* 结论：尽管存在性能开销，但性能表现并非灾难性的，仍然可用。

性能对比图，显示了确定性模式与基线模式在吞吐量、首次Token延迟和Token间延迟方面的差异。
性能对比图，显示了确定性模式与基线模式在吞吐量、首次Token延迟和Token间延迟方面的差异。

真正的在线策略强化学习（True on-policy RL）
* 实验内容：研究人员指出，训练和推理之间的数值差异会隐式地将在线策略RL（on-policy RL）转变为离线策略RL（off-policy RL）。本实验在RLVR设置下，比较了三种情况：1) 不使用离线策略校正；2) 使用重要性加权进行校正；3) 使用确定性推理实现真正的在线策略。
* 实验结果：
1. 不使用校正时，奖励在训练中途崩溃。
2. 使用重要性加权校正后，训练可以平稳进行。
3. 通过在采样器和训练器之间实现逐比特相同的结果，实现了完全的在线策略，训练同样平稳。
KL散度图显示，无校正时KL散度在奖励崩溃时飙升；有校正时KL散度维持在0.001左右；而在“真正的在线策略RL”中，KL散度始终为0，表明训练策略和采样策略之间没有偏差。
* 结论：确定性推理是实现真正的在线策略RL的关键，能够消除因数值不一致性引入的策略偏差。

A5 结论

现代软件系统包含多层抽象。在机器学习中，当我们遇到不确定性和细微的数值差异时，很容易选择忽视它们。毕竟，我们的系统已经是“概率性”的，多一点不确定性似乎无伤大雅。然而，本文反对这种失败主义观点。通过一些努力，我们能够理解不确定性的根本原因并解决它们。作者希望这篇博文能为社区提供一个关于如何解决推理系统中不确定性的坚实理解，并激励其他人去全面理解他们自己的系统。