批次不变性才是真正根源

大语言模型推理结果的不可复现性长期困扰着研究者和工程师。即使将采样温度设为零，同一问题在不同时刻仍可能得到不同答案。一种广为流传的解释认为，GPU 的高度并行化使得浮点数加法顺序随核心完成先后而改变，进而导致输出差异。然而 Thinking Machines 团队通过实验发现，对同一组数据反复执行矩阵乘法，每次结果在比特层面完全一致，这说明并发执行本身并非不确定性的直接来源。

真正的根源在于浮点数的非结合性与批次大小的耦合。浮点加法在操作数量级差异较大时会丢失精度，导致不同加法顺序产生不同结果。LLM 推理服务器的前向传播在给定完全相同输入时是确定性的，但服务器的实时负载决定了每次前向传播的批次大小。批次大小的变化会改变矩阵乘法、RMSNorm 和注意力机制内部的归约策略，从而使单个用户请求的输出依赖于同时在线的其他用户数量——这正是用户视角下不确定性的真实来源。

研究团队将问题分解为三类需要归约操作的算子：RMSNorm、矩阵乘法和注意力机制，并逐一分析批次不变性的实现路径。对于 RMSNorm，数据并行策略将每个批次元素分配给独立核心处理，批次大小增大时不影响归约顺序；批次极小时若引入分裂归约则会破坏批次不变性，最简单的处理方式是对小批次场景不做额外优化。矩阵乘法同理，固定使用单一内核配置可保证批次不变性，实测性能损失约 20%，在 LLM 推理场景中尚属可接受范围。

注意力机制的批次不变性实现难度最高。解码阶段查询序列极短，若不沿 KV 维度分裂归约则无法充分利用 GPU 算力，而常见的动态分裂策略（如 FlashInfer 的均衡调度）会使归约顺序随当前处理的查询 token 数量变化。研究团队提出固定分裂块大小而非固定分裂数量的策略：无论处理多少 token，每个分裂块的大小保持不变，从而保证归约顺序的一致性。此外，还需在注意力内核执行前统一更新 KV 缓存布局，消除缓存与当前 token 分开处理带来的边界条件差异。

研究团队基于 vLLM 的 FlexAttention 后端实现了上述批次不变性内核，并开源了相关代码。实验使用 Qwen3-235B 模型，以温度零对同一提示生成 1000 次输出：默认配置下产生 80 种不同结果，启用批次不变性内核后 1000 次输出完全一致。性能方面，未优化的确定性 vLLM 耗时约为默认配置的两倍，引入改进注意力内核后可降至约 1.6 倍，仍有较大优化空间。更重要的是，推理确定性为强化学习提供了新的可能：当采样器与训练器在比特层面完全对齐时，KL 散度恒为零，无需重要性加权即可稳定训练，避免了因数值不匹配导致的奖励崩溃。

要点

• LLM 推理不确定性的真实根源不是 GPU 并发竞争，而是批次大小变化导致关键算子的归约顺序改变，即缺乏批次不变性
• 通过为 RMSNorm、矩阵乘法和注意力机制分别设计批次不变性内核，可在 vLLM 上实现 1000 次请求结果完全一致
• 固定注意力分裂块大小而非固定分裂数量，是实现解码阶段批次不变性的关键工程决策
• 确定性推理使采样器与训练器逐位对齐成为可能，从而实现零 KL 散度的真正在线策略强化学习，无需重要性加权即可稳定训练
• 当前批次不变性内核的性能损失约为 1.6 倍，主要来自 FlexAttention 集成尚未充分优化，仍有较大提升空间