LLM 推理的非确定性：根因分析与 Batch Invariance 解决方案

2025-12-24 · Qi Lu · Views:

问题定义

在 LLM 推理服务中，相同的输入理应产生相同的输出。然而实际观测表明，即使采用贪婪解码（temperature=0），输出仍存在不确定性。以下实验数据来自 Thinking Machines Lab1：

这一现象与贪婪解码的数学定义矛盾：$\hat{y}_t = \arg\max_v p(v \mid y_{<t}, x)$ 应当是确定性的。

本文的目标是：(1) 定位非确定性的根本来源；(2) 设计可工程化部署的解决方案。

假说检验：并发浮点运算 vs. Batch Size 变化

假说 1：GPU 并发导致的浮点非关联性

主流假说认为，GPU 的并行计算引入了不确定性。其逻辑链条如下：

1. 浮点加法不满足结合律：$(a + b) + c \neq a + (b + c)$
2. GPU 并行 reduction 的执行顺序取决于线程调度
3. 不同的执行顺序产生不同的累加结果
4. 微小差异经 argmax 放大后改变输出 token

这一假说的典型支持证据是 CUDA 的 atomicAdd 操作的非确定性。

假说 1 的问题

Thinking Machines Lab1 的分析指出，现代 Transformer 推理中的核心操作并不依赖原子操作：

• GEMM：使用 cuBLAS 的分块矩阵乘法，reduction 树结构固定
• LayerNorm/RMSNorm：标准实现使用确定性的 warp-level reduction
• Attention：FlashAttention 的 tiled 实现同样采用固定 reduction 顺序

实验验证：在单 GPU、固定 batch size 条件下，同一输入的多次推理输出完全一致。这排除了线程级非确定性作为主要来源。

假说 2：Batch Size 变化导致数值路径分歧

真正的问题在于：kernel 的数值输出是 batch size 的函数。

当推理服务采用动态 batching 时，同一请求在不同时刻可能被分配到不同大小的 batch 中。每个 kernel 的 tiling 策略、reduction 分块方式都可能随 batch size 变化，导致数值结果不同。

Batch Invariance 的数学定义

设 kernel $K$ 作用于输入张量 $X \in \mathbb{R}^{B \times N \times D}$，输出 $Y = K(X)$。

Batch Invariance 要求：对于任意 batch 中的样本 $x_i$，其输出 $y_i$ 仅取决于 $x_i$ 本身，而与 batch 中其他样本无关：

其中 $K’$ 是等价的单样本 kernel。

实际上，这一性质在标准 kernel 实现中并不成立。原因在于 tiling 策略的 batch 依赖性。

关键 Kernel 的 Batch Variance 分析

RMSNorm

RMSNorm 计算如下：

关键步骤是 reduction：$\sum_{i=1}^{D} x_i^2$。

标准 CUDA 实现通常采用 tree reduction：

问题所在：当 batch size 改变时，CUDA kernel 可能选择不同的 block size 和 grid 配置。不同配置下的 reduction 分块方式不同：

• Batch size = 1：可能使用 256 threads/block，reduction 分 4 轮完成
• Batch size = 8：可能使用 128 threads/block，reduction 分 8 轮完成

不同的分块方式产生不同的中间舍入误差，最终结果在 ULP（Unit in the Last Place）级别存在差异。

矩阵乘法

GEMM 的标准分块实现：$C = AB$，其中 $A \in \mathbb{R}^{M \times K}$，$B \in \mathbb{R}^{K \times N}$。

cuBLAS 根据矩阵尺寸和 GPU 架构选择最优 tiling 配置。例如：

当 batch size 改变时，等效的矩阵尺寸 $M’ = B \times M$ 发生变化，触发不同的 tiling 选择，导致 K 维 reduction 的分块方式不同。

数学表达：设 K 维被分为 $P$ 个块，每块大小 $k_p$：

浮点加法的顺序取决于 $P$ 和各块的边界划分。不同的 $P$ 产生不同的结果。

Attention

FlashAttention 的核心是 分块计算 + Online Softmax。Softmax 的增量更新公式：

\(m^{new} = \max(m^{old}, m^{(j)})\) \(\ell^{new} = e^{m^{old} - m^{new}} \ell^{old} + e^{m^{(j)} - m^{new}} \ell^{(j)}\) \(O^{new} = \frac{e^{m^{old} - m^{new}} \ell^{old} \cdot O^{old} + e^{m^{(j)} - m^{new}} \ell^{(j)} \cdot O^{(j)}}{\ell^{new}}\)

关键参数是 KV split size：将 KV cache 分成多少块进行增量计算。

在 decoding 阶段，FlashAttention 根据 KV cache 长度和 GPU 配置动态选择 split size。不同的 split size 导致不同数量的 rescaling 操作，累积不同的舍入误差。

误差传播：从 ULP 到 Token 分歧

单个 kernel 的数值差异通常在 $10^{-6}$ 到 $10^{-8}$ 量级。如何导致 token 级别的分歧？

误差累积模型

设 Transformer 有 $L$ 层，每层的相对误差为 $\epsilon$。最终 logits 的相对误差约为：

对于 $L = 80$（如 Llama-70B）、$\epsilon = 10^{-7}$：

Argmax 的脆弱性

当两个 token 的概率接近时，微小的 logits 差异足以翻转 argmax：

若 $|\Delta \ell| < \epsilon_{total}$，则结果不稳定。

实际观测表明，在 greedy decoding 中，约 1-5% 的 token 位置存在这种”脆弱”状态。一旦发生分歧，后续生成完全不同，呈现 蝴蝶效应。

解决方案：Batch-Invariant Kernel 设计

设计原则

实现 batch invariance 的核心原则：

1. 固定 tiling 配置：无论输入尺寸如何，使用固定的 block size 和 grid 配置
2. 固定 reduction 顺序：确保 reduction 树结构与 batch size 无关
3. 固定 split size：Attention 中使用固定的 KV split，不随序列长度动态调整

RMSNorm 的 Batch-Invariant 实现

关键修改：强制使用固定的 reduction 分块。

# 标准实现（batch-variant）
def rmsnorm_standard(x, gamma, eps):
    # reduction 分块由 CUDA runtime 决定
    rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + eps)
    return x / rms * gamma

# Batch-invariant 实现
def rmsnorm_batch_invariant(x, gamma, eps, fixed_block_size=256):
    # 强制固定分块
    D = x.shape[-1]
    num_blocks = (D + fixed_block_size - 1) // fixed_block_size

    # 分块累加，固定顺序
    sq_sum = 0.0
    for i in range(num_blocks):
        start = i * fixed_block_size
        end = min(start + fixed_block_size, D)
        sq_sum = sq_sum + torch.sum(x[..., start:end] ** 2, dim=-1, keepdim=True)

    rms = torch.sqrt(sq_sum / D + eps)
    return x / rms * gamma

实际的 CUDA 实现需要确保 warp-level reduction 的顺序固定，通常通过显式的 __shfl_down_sync 序列实现。

MatMul 的 Batch-Invariant 实现

cuBLAS 的自动 tuning 会根据矩阵尺寸选择不同 kernel。解决方案：

1. 禁用自动 tuning：强制使用固定的 GEMM kernel
2. Padding 到固定尺寸：将矩阵 pad 到 $2^n$ 大小，确保相同的 tiling

# 使用 torch.Library 替换默认 kernel
import batch_invariant_ops

# 自动替换所有 matmul 为 batch-invariant 版本
# 内部实现固定 tile size = 128x128

Attention 的 Batch-Invariant 实现

关键：固定 KV split size。

# FlashAttention batch-invariant 配置
flash_attn_func(
    q, k, v,
    deterministic=True,      # 启用确定性模式
    fixed_split_kv=64        # 固定 KV split 为 64
)

SGLang 的实现支持多种 attention 后端的 batch-invariant 模式：

多 GPU 场景：AllReduce 的确定性

在 Tensor Parallelism 中，AllReduce 操作同样引入非确定性。NCCL 的默认实现使用 ring-based 或 tree-based 算法，reduction 顺序取决于 GPU 间通信延迟。

确定性 AllReduce

解决方案：使用固定顺序的 reduce-scatter + all-gather：

SGLang 实现了 deterministic tensor parallelism，确保多 GPU 场景下的完全可复现。

MoE 模型的额外挑战

Mixture-of-Experts（MoE）架构引入了 Dense 模型不存在的非确定性来源。以 Qwen3-235B-A22B、DeepSeek-V3 等模型为代表的 MoE 架构，其稀疏激活特性使得确定性推理面临额外挑战。

Token Routing 的非确定性

MoE 的核心是 门控网络（Gating Network），决定每个 token 被路由到哪些 expert：

问题在于：当多个 expert 的门控分数接近时，微小的数值扰动可能改变 TopK 的选择结果。

Expert Capacity 与 Token Dropping

为保证计算效率，MoE 通常设置 Expert Capacity：每个 expert 能处理的最大 token 数。当某个 expert 接收的 token 超过容量时，多余的 token 被 丢弃（dropped），直接通过残差连接传递到下一层。

Token dropping 的非确定性来源：

1. Batch 组成变化：不同 batch 中的 token 分布不同，同一 token 可能在某些 batch 中被处理，在另一些中被丢弃
2. 路由顺序依赖：当 capacity 接近饱和时，先到达的 token 被处理，后到达的被丢弃
3. 负载不均衡：热门 expert 更容易触发 dropping

训练-推理不一致

研究发现，MoE 模型在训练和推理阶段的路由行为存在显著差异：

“Even under identical conditions, the routing framework can yield divergent expert selections across repeated forward passes.”

这种不一致在 RL 训练中尤为严重：

• On-policy 要求：RL 训练要求 rollout 与训练使用相同的策略
• 路由漂移：推理时的路由决策可能偏离训练时的分布
• 奖励噪声：非确定性路由引入难以追踪的奖励波动

MoE 确定性推理的解决方案

1. 固定路由阈值

避免在门控分数接近时产生不稳定决策：

def deterministic_topk(scores, k, margin=1e-5):
    # 当分数差异小于 margin 时，使用固定的 tie-breaking 规则
    sorted_scores, indices = torch.sort(scores, descending=True)

    # 检测 tie 情况
    for i in range(k-1, len(sorted_scores)-1):
        if sorted_scores[i] - sorted_scores[i+1] < margin:
            # 使用 expert index 作为 tie-breaker
            pass

    return indices[:k]

2. 取消 Token Dropping

以计算效率换取确定性：

# 设置足够大的 capacity factor，确保不发生 dropping
config.moe_capacity_factor = 2.0  # 或更高
config.moe_drop_tokens = False

3. Soft MoE

使用软路由替代硬路由，每个 token 以加权方式分配给所有 expert：

Soft MoE 消除了离散的路由决策，但计算开销更高。

对 RL 训练的特殊影响

MoE 模型的 RL 训练面临双重挑战：

1. 数值非确定性：前述的 batch variance 问题
2. 结构非确定性：路由决策的不稳定性

研究表明，MoE RL 训练的不稳定性部分源于路由分布的漂移：

“The routing distribution has been identified as a pivotal factor contributing to the instability of MoE RL.”

实验数据显示，约 10% 的 router 在训练和推理间选择不同的 expert，94% 的 token 至少在一层选择了不同的 expert。

Routing Replay：R2 与 R3

为解决训练-推理路由不一致问题，研究者提出了 Routing Replay 机制，核心思想是在训练时重放推理阶段的路由决策。

R2：Vanilla Routing Replay

定义：重用 训练系统 在 rollout 阶段选择的 expert。

Rollout (Training System) → 记录路由决策 → Training 时重放

• 优点：实现简单，开销较小
• 缺点：当 off-policy 程度较大时效果下降

R3：Rollout Routing Replay

定义：重用 推理系统 在 rollout 阶段选择的 expert。

Rollout (Inference System) → 记录路由决策 → Training 时重放

R3 强制约束：在 rollout 阶段激活的特定 expert 必须在训练反向传播时严格重用。

选择 R2 还是 R3？

实现细节

R3 可与 KV Cache 集成：

# 对于每层每个 token prefix，存储对应的路由 mask
routing_cache = {}

def forward_with_replay(x, layer_idx, prefix_hash):
    if prefix_hash in routing_cache:
        # 命中缓存，重用路由决策
        mask = routing_cache[prefix_hash]
    else:
        # 计算新路由
        mask = compute_routing(x)
        routing_cache[prefix_hash] = mask

    # 使用 mask 进行 softmax gating，但不阻断梯度流
    return apply_routing(x, mask, allow_grad=True)

关键点：重放的 mask 用于 softmax gating 计算，但不阻止梯度流经标准 router 权重，确保 router 仍可训练。

其他解决方案

RSPO（Router-Shift Policy Optimization）

不同于 R2/R3 的硬约束，RSPO 采用软调整机制：

• 计算当前策略与旧策略间的 router shift ratio
• 量化每个 token 的路由偏移程度
• 自适应重加权更新，限制过大更新同时保持 router 灵活性

冻结 Router

最简单的方案是冻结 router 参数，但会损害模型适应性，通常不推荐。

性能分析

Batch-invariant kernels 的主要开销来源：

1. 放弃动态优化：无法使用针对特定尺寸的最优 kernel
2. 固定 tiling 的低效：小矩阵使用大 tile 造成资源浪费
3. 额外的同步开销：确保固定执行顺序需要更多同步点

基准测试

CUDA Graphs 通过预编译执行图，消除了 kernel launch 开销，显著减少了确定性模式的性能损失。

延迟分布

确定性模式的另一个优势是 延迟方差降低：

标准模式:   P50=45ms, P99=120ms, stddev=25ms
确定性模式: P50=52ms, P99=58ms,  stddev=3ms

对于延迟敏感的应用，更低的方差可能比更低的均值更重要。

应用场景

RL 训练的可复现性

在 RLHF/RLVR 训练中，策略的 rollout 需要与训练步骤精确对应。非确定性推理导致：

1. 调试困难：无法复现特定的失败案例
2. 训练不稳定：相同 checkpoint 在不同机器上表现不一致
3. 实验不可比较：无法区分算法改进和随机波动

SGLang 与 slime 的合作实现了 100% 可复现的 RL 训练：

“Taking this deterministic inference capability further, SGLang collaborated with the slime team to unlock 100% reproducible RL training.”

安全与合规

对于需要审计的 AI 系统，确定性行为是基本要求：

• 医疗诊断
• 金融决策
• 法律分析

非确定性意味着无法追溯和解释特定输出的来源。

模型调试

当模型出现异常输出时，确定性推理允许：

1. 精确复现问题
2. 逐层检查中间状态
3. 二分定位问题来源

相关工作

LayerCast（NeurIPS 2025）

Yuan et al. 提出了另一种思路：通过提高数值精度减少误差累积。

• 方法：权重以 FP16/BF16 存储，计算时转换为 FP32
• 优势：不修改 kernel 实现
• 局限：无法完全消除 batch variance，只是降低其影响

实验显示，LayerCast 将 DeepSeek-R1-Distill-Qwen-7B 的准确率波动从 9% 降低到 2%，但仍非完全确定。

OpenAI seed 参数

OpenAI API 提供 seed 参数以提高可复现性：

“If specified, our system will make a best effort to sample deterministically…”

但官方明确表示 不保证确定性，原因包括：

1. 后端模型更新
2. 负载均衡到不同硬件
3. 系统配置变化

system_fingerprint 字段用于追踪配置变化，但无法保证相同 fingerprint 下的完全确定性。

总结

LLM 推理的非确定性源于 kernel 实现对 batch size 的敏感性，而非 GPU 的并发特性。通过设计 batch-invariant kernels，可以在工程层面实现完全确定的推理。

Dense 模型的关键技术点：

1. 固定 RMSNorm 的 reduction 分块
2. 固定 MatMul 的 tiling 配置
3. 固定 Attention 的 KV split size
4. 多 GPU 场景使用确定性 AllReduce

MoE 模型的额外挑战：

1. Token routing 的离散决策在门控分数接近时不稳定
2. Expert capacity 和 token dropping 引入 batch 依赖的非确定性
3. 训练-推理路由不一致需要 Routing Replay（R2/R3）机制

性能代价约 30-60%，但在 RL 训练、模型调试、安全审计等场景中是必要的投入。

参考文献

[1] Thinking Machines Lab. Defeating Nondeterminism in LLM Inference. 2025.

[2] Yuan, J. et al. Understanding and Mitigating Numerical Sources of Nondeterminism in LLM Inference. NeurIPS 2025 (Oral).

[3] LMSYS Org. Towards Deterministic Inference in SGLang and Reproducible RL Training. 2025.

[4] Thinking Machines Lab. batch_invariant_ops. GitHub.

[5] vLLM Documentation. Batch Invariance.

[6] Dao, T. et al. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023.

[7] Ma, W. et al. Stabilizing MoE Reinforcement Learning by Aligning Training and Inference Routers. 2025.

[8] Towards Stable and Effective Reinforcement Learning for Mixture-of-Experts. arXiv 2025.