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

2025-12-24 · Qi Lu · Views:

问题定义

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

这一现象与贪婪解码的数学定义矛盾：$\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 Lab 的分析指出，现代 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：

实际观测表明，在 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.