@jessiedong_: split-K matrix multiplications are known for giving different answers even when the code and inputs are the same. but h…

split-K matrix multiplications are known for giving different answers even when the code and inputs are the same. but how different are they actually?

I couldn’t find anyone who’s measured how this changes with block layouts and split counts! when running the same atomic split-K kernel 200 times on a B200, when blocks for each output were far apart in the launch order, every answer matched.

when I moved the split index so those blocks were next to each other, all 4096 outputs changed at least once with 32 splits.

1/n

in gpu kernels, split-K divides one matrix multiplication so more blocks can work on it at the same time. this can make smaller matrix multiplications faster by using more of the gpu.

with atomic split-K, every block adds its partial answer into the same output when it finishes. this is where the order can change.

2/n

why does the order matter?

the gpu can’t store every number exactly, so it sometimes has to round after each addition.

if you add the same numbers in a different order, the rounding can happen in different places. so you can end up with a slightly different answer (though none of the numbers changed).

3/n

there were 4 parts:

• to measure differences: the same calculation was repeated while changing K, the number of splits and the number format
• to see why the answers changed: order that blocks added their answers was recorded
• to see if a model could notice: outputs were used as token scores
• to see the speed difference between repeatable and atomic versions: regular, atomic, and fixed-order versions were timed

4/n

in the first part, the same custom atomic split-K kernel was run 500 times on a B200 using the same inputs.

the test covered four K sizes, six split counts and fp32, fp16 and bf16 inputs. across all 96 setups, every output matched on every run!

5/n

the second part measured whether the blocks were adding their answers in a different order.

CUDA’s atomicAdd function returns the old output value before adding the new one. saving that value for every block shows which answer was added first, second and so on.

6/n

the first layout grouped the work by K split. all 4096 output blocks for the first piece of K came before the blocks for the second piece, and so on.

this kept the blocks for one output far apart. on the B200, their answers reached the output in the same order on every run, which explains why nothing changed.

7/n

the second layout put K-split blocks for each output next to each other.

the blocks could now finish in different orders. across 200 runs, every output had its partial answers added in a different order at least once.

8/n

as the number of splits increased, more outputs changed:

2 splits: 0/4096 4 splits: 3015/4096 8 splits: 4092/4096 32 splits: 4096/4096

with the first layout, 0/4096 outputs changed at every split count

9/n

why did 2 splits still give the same output?

there were only two partial answers, and adding a+b gives the same result as b+a. with 4 or more partial answers, changing the order also changes which values get added and rounded first. this is when outputs started to differ.

10/n

the largest differences compared with the original output were:

4 splits: 0.005% 8 splits: 0.013% 32 splits: 0.031%

something to think about is where these differences could matter:

• during normal inference, probably only when two possible next tokens have almost the same score
• during RL, the chosen token can stay the same but the number used to train on it can change
• during testing, matching answers make it easier to tell if a code change caused something to change

11/n

lastly, there were 1) small changes to token probabilities and 2) almost no extra time from using a fixed order

• token scores: repeating same atomic split-K version gave same probabilities every time. switching between no split-K and split-K changed log probabilities by up to 1.87e-6

• speed: fixed-order version was ~0–2% slower than atomic version in this test. split-K was slower for smaller K and only ~5% faster in its best result

12/n

how work is arranged across a gpu is important to think about. thinking back to when I started thinking about this more carefully, it was probably the Thinking Machines experiments in september ’25 on batch invariance, which talked a lot about repeating the same matrix multiplication giving the same answer, but the nuances of this when you change how the work is divided! script: http://github.com/jessiedong01/splitk-nondeterminism…

13/13