原创 小鹿 2025-09-11 17:16 北京
博客链接：https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/这篇博客是公司的新栏目 “Connectionism”（连接主义） 的首秀，TML 公司表示他们将频繁分享研究与代码：要做真正的 Open 的 AI 呗，在阴阳谁我不说 ～TML 联合创始人&著名技术博主 Lilian Weng 翁荔（也是 OpenAI 前研究与安全副总裁 VP） 随后转推，暗示了“Connection Machine”的存在。（不知道这会不会是他们未来产品的名字！）而这篇开山之作文章的核心作者，正是刚从 Meta 离职的 PyTorch 核心开发者 Horace He。他要挑战的，正是那个让无数开发者头疼的难题：为什么与大模型的对话，总像在“开盲盒”？问题定位但是理论上，所有开发者都知道的另一个事实，就是将温度系数（temperature）设为零，彻底关掉随机性。而这时的输出，本应像“1+1=2”一样，是绝对稳定可预测的。它还是会变。这个“通病”让模型评测变得困难，让调试过程如同噩梦，更让那些需要绝对确定性的下游应用（如强化学习、代码生成）也很头疼。在此之前，社区有一个广为流传的关于这个“通病”的主流假说—“并发 + 浮点”理论。懂计算机 float 运算的小伙伴，这其实很好理解。第一，是浮点计算的非结合性导致的，由于计算机存在舍入误差，数学上 (a+b)+c = a+(b+c) 的绝对相等，在浮点数世界里会变成约等，而计算顺序的改变，会带来结果上“位级别”的微小差异。第二，是GPU并行计算决定的：为了追求极致速度，GPU 会同时派出成百上千个核心（kernel）去执行求和等归约操作。谁先算完、谁后算完，这个计算的顺序在每次运行时都是不可预测的。当上面这两个随机不可控撞在一起，偏差就出现了。这一点点偏差，就足以让模型在某个关键的“岔路口”做出不同选择，从而生成不同的输出结果。不过，Thinking Machines Lab 发布的《克服 LLM 推理中的不确定性》论文博客提出了质疑。他们认为，问题不在“有没有并行/浮点”。看看他们是怎么论证的？首先研究团队做了一个简单的实验：在 GPU 上，反复对相同的数据做同一个 matmul （矩阵乘法），结果在位级别上是完全确定的。如果之前主流假说成立，即并发计算必然导致随机性，那么这里也应该出现不一致。研究团队另外提到，在典型的 LLM 前向传播中，为了性能与稳定性，所采用的底层计算内核（Kernels）已经主动避免了使用那些会引入 running-time 不确定性的技术，例如跨线程块的原子加法。原子加法操作虽然能保证并发计算结果的完整性，但硬件并不保证这些计算的累加顺序。这种顺序的随机性是“原子操作”与生俱来的特性，一旦算法依赖它，几乎必然会导致非确定性。
但是，目前 LLM 前向传播的核心环节（如矩阵乘法、层归一化）并不涉及这类原子操作。这意味着，LLM 的前向过程本质上是 running-time 确定的，即对于固定的输入，无论运行多少次，它都会产生完全相同的位级别输出。所以 TML 的研究团队，认为不确定性来自什么时候归约顺序会被改写。所以论文将矛头指向了一个更高阶的系统设计缺陷：批次不变性 (Batch Invariance) 的缺失。“批次不变性”是指：
一个数据样本的计算结果，不应受到其所在处理批次（batch）的大小或批次中其他样本内容的影响。那通俗地说：
你向 AI 提问，得到的结果，理论上不应该因为和你一同被处理的还有 3 个人还是 30 个人（即批次大小）而有所不同。但在现实的推理服务器上，正是推理 inference 阶段的存在的问题，为了效率这个原则被打破了：服务器负载是随机的，你永远不知道在你提问的那一刻，服务器上还有多少并发请求。动态批处理，为了让 GPU 不“挨饿”，服务器会把短时间内收到的多个请求“打包”成一个批次（Batch）再处理。批次大小（Batch Size）变得随机，你的请求这次可能被分入一个大小为 4 的批次，下次可能是大小为 16 的批次。计算策略为性能而动态切换，底层的计算内核（Kernel）为了在任何批次大小下都跑出最快速度，被设计得极其“智能”：它会根据批次大小，动态选择最优的计算策略。解决问题定位到问题后，研究团队认为解决方案也就清晰了：强制计算内核（Kernels）使用固定的策略，牺牲一部分针对特定情况的极致性能优化，以换取计算路径的绝对一致性。论文展示了如何对 Transformer 模型的三个核心模组进行改造：RMSNormRMSNorm 的改造相对直接，它的标准并行化策略是“数据并行”，即为批次中的每个序列（sequence）分配一个 GPU 的计算单元（线程块/SM）。这种策略在批次较大时（例如，批次大小超过 GPU 核心数）效率很高。但当批次很小时，大量计算单元会处于空闲状态，造成严重的性能浪费。为了解决小批次下的效率问题，高性能计算内核会采取一种“自适应”策略：它会动态检测到核心空闲，并立刻改变并行模式，不再一个核心处理一个序列，而是用多个核心协同处理一个序列。这种“分裂式规约”为了合并结果，引入了额外的跨核心通信和求和步骤，彻底改变了原始的计算顺序，从而破坏了批次不变性。最直接的解决方案就是禁用这种自适应并行策略，开发者必须选择并固定一种对所有批次大小都适用的归约策略。矩阵乘法矩阵乘法（Matmul）的非不变性来源更为复杂：1. Split-K 策略： 与 RMSNorm 类似，对于某些矩阵形状（特别是当批次维度 M 和 N 较小，而需要累加的 K 维度较大时），为了创造更多并行任务让 GPU“忙起来”，一种常见的优化是“Split-K”。它将 K 维度切分成多块，分配给不同核心并行计算，最后再将各部分结果相加。这种“先分后总”的模式，从根本上改变了加法顺序。2. 张量核心指令的动态选择： 现代 GPU 库中，包含一个庞大的、针对不同硬件和数据类型高度优化的计算函数库。在运行时，一个“启发式调度器”会根据输入张量的具体形状，动态地选择一个它认为最快的函数实现。批次大小的微小改变，就可能导致调度器选择一个完全不同的底层函数，而不同函数的内部计算顺序可能天差地别。解决方案是放弃这种动态调度和优化。开发者为模型中遇到的所有矩阵乘法形状，编译并指定一个固定的计算内核配置（固定的 tile 大小、固定的并行策略，并明确禁用 Split-K）。这牺牲了针对特定形状的极致性能优化，以换取一条稳定不变的计算路径。但是幸运的是，在 LLM 推理场景下，模型维度通常很大，禁用 Split-K 等优化的性能损失通常在可接受范围内。注意力机制注意力机制的确定性改造更困难，因为它引入了全新的复杂维度：跨序列维度归约： 注意力的核心是计算查询（Query）与键（Key）的点积，这涉及在序列长度维度上的归约，而序列长度本身是动态变化的。与推理优化的深度耦合： 在 vLLM 等现代推理引擎中，为了高效管理内存，使用了分页 KV 缓存（Paged KV Cache）等复杂技术，注意力计算必须与这些动态变化的内存布局进行交互。因此研究团队希望的是，对于任意一个给定的查询 token，无论它是在预填充阶段与数百个其他 token 一同处理，还是在解码阶段单独处理，其最终的注意力计算结果必须完全相同。于是，研究团队提出了一个两阶段的确定性解决方案：确保一致的内存布局： 在调用注意力计算函数之前，推理引擎必须先更新 KV 缓存和相应的页表，确保无论处理上下文如何（预填充或解码），键（Key）和值（Value）张量在内存中的布局都是一致的。这从源头上消除了因内存访问模式不同而引入的差异。针对 Split-KV 的“固定拆分大小”策略： 在解码阶段，查询长度通常为 1，此时为了保持 GPU 繁忙，几乎必须沿 KV 序列长度维度进行并行化（即“Split-KV”）。
传统做法（非确定性）：根据当前批次的总工作量，动态地将 KV 序列切分成最合适的块数，以期最大化并行度。这种做法使得切分方式依赖于批次大小，从而破坏了不变性。而论文提出的是新做法是采用“固定拆分大小”（fixed split-size）策略，不再固定切分的数量，而是固定每个切分块的大小（例如，固定为 256 个 token）。一个长度为 1000 的 KV 序列，将被确定性地切分为 3 个大小为 256 的块和 1 个大小为 232 的块，确保归约的结构完全独立于批次中的其他请求，从而完美地保持了批次不变性。实验论证论文设计了一系列实验，极其直观地展示了其解决方案的有效性，并量化了实现确定性所需付出的代价。生成结果对比为了让读者直观感受“非确定性”的严重程度，研究团队进行了一项生成实验。实验场景： 使用 Qwen 模型，将 temperature 设为 0，对提示词“Tell me about Richard Feynman”进行 1000 次采样。测试模型：标准 vLLM（非确定性）：1000 次运行，竟产生了 80 个不同的输出版本。大部分（992 次）结果在前 102 个词元后生成了“Queens, New York”，而少数（8 次）结果则生成了“New York City”。这清晰地表明，差异并非随机噪声，而是由系统非不变性触发的、微小但具有决定性的计算路径分叉。确定性 vLLM（改造后）： 全部 1000 次采样，生成了完全相同、逐位一致的唯一结果。性能代价那换取这种 100% 的可靠性，代价是什么？研究团队设置了一个 API 服务器，其中一台 GPU 运行 Qwen-3-8B，并请求 1000 个序列，输出长度在 90 到 110 之间。
