@jessiedong_: split-K 矩阵乘法即使在代码和输入相同的情况下也以给出不同的答案而闻名。但 h…

split-K矩阵乘法以即使代码和输入相同也会产生不同结果而闻名。但它们究竟有多大差异？

我找不到有人测量过这如何随分块布局和分割数量而变化！当在B200上对同一个原子split-K内核运行200次时，如果每个输出的分块在启动顺序中相距很远，所有答案都匹配。

当我移动分割索引使这些分块彼此相邻时，32次分割下所有4096个输出至少改变了一次。

1/n

在GPU内核中，split-K将一个矩阵乘法分割开来，以便更多分块可以同时处理。这可以通过使用更多GPU资源来加速较小的矩阵乘法。

在原子split-K中，每个分块在完成时将其部分答案添加到同一个输出中。这就是顺序可能改变的地方。

2/n

为什么顺序很重要？

GPU无法精确存储每个数字，因此有时每次加法后都必须进行舍入。

如果你以不同的顺序添加相同的数字，舍入可能发生在不同的地方。因此你最终可能得到一个略微不同的答案（尽管没有数字发生改变）。

3/n

研究包含四个部分：

• 测量差异：在改变K值、分割数量和数字格式的同时重复相同计算
• 查看答案为何改变：记录分块添加答案的顺序
• 查看模型是否能察觉：将输出用作token分数
• 查看可重复版本与原子版本之间的速度差异：对常规、原子和固定顺序版本进行计时

4/n

在第一部分中，使用相同的自定义原子split-K内核，在B200上用相同输入运行500次。

测试涵盖四种K值大小、六种分割数量以及fp32、fp16和bf16输入。在所有96种设置中，每次运行的每个输出都匹配！

5/n

第二部分测量了分块是否以不同的顺序添加它们的答案。

CUDA的atomicAdd函数在添加新值之前返回旧输出值。为每个分块保存该值可以显示哪个答案是第一个、第二个添加的，等等。

6/n

第一种布局按K分割分组工作。第一段K的所有4096个输出分块都在第二段的分块之前，依此类推。

这使得一个输出的分块相距很远。在B200上，它们的答案每次运行都以相同的顺序到达输出，这解释了为什么没有任何改变。

7/n

第二种布局将每个输出的K分割分块彼此相邻放置。

分块现在可以以不同的顺序完成。在200次运行中，每个输出的部分答案至少有一次是以不同的顺序添加的。

8/n

随着分割数量增加，更多输出发生了改变： 2次分割：0/4096 4次分割：3015/4096 8次分割：4092/4096 32次分割：4096/4096

使用第一种布局时，每个分割数量下都是0/4096输出发生改变

9/n

为什么2次分割仍然给出相同输出？

只有两个部分答案，而a+b与b+a的结果相同。对于4个或更多部分答案，改变顺序也会改变哪些值首先被添加和舍入。这时输出开始出现差异。

10/n

与原始输出相比的最大差异为： 4次分割：0.005% 8次分割：0.013% 32次分割：0.031%

值得思考的是这些差异可能在哪些情况下重要：

• 在正常推理期间，可能只在两个可能的下一个token分数几乎相同时
• 在强化学习期间，选择的token可能保持不变，但用于训练它的数字可能改变
• 在测试期间，匹配的答案更容易判断代码更改是否导致某些变化

11/n

最后，有两点：1) token概率有微小变化，2) 使用固定顺序几乎没有额外时间开销

• token分数：重复相同的原子split-K版本每次给出相同的概率。在不使用split-K和使用split-K之间切换使对数概率最多改变1.87e-6
• 速度：在此测试中，固定顺序版本比原子版本慢约0-2%。split-K在较小时更慢，在其最佳结果中仅快约5%

12/n

思考如何在GPU上安排工作很重要。回想我开始更认真地思考这个问题的时候，可能是在2025年9月Thinking Machines进行的批处理不变性实验，该实验大量讨论了重复相同矩阵乘法得到相同答案，但当你改变工作分配方式时，这种情况的微妙之处！ 脚本：http://github.com/jessiedong01/splitk-nondeterminism…