Thinking Machines Lab’s batch-invariance problem, formalized

Robert Joseph George · TorchLean Verified Examples

We give a model the same prompt twice, but putting one request in a different batch can change its answer. Even with sampling turned off! The Thinking Machines Lab post on batch-invariant inference explains how this happens: the other requests can affect how the backend groups arithmetic operations, and floating-point addition depends on that grouping. I wanted to work through that calculation in TorchLean and FloatLib.

Let's start with three numbers whose sum changes when we move the parentheses. We'll evaluate both orders in Lean, then prove conditions under which changing the batch preserves the scores or the chosen token.

Temperature zero removes randomness from the token-selection rule; it doesn't fix the arithmetic that produces the scores. Greedy decoding chooses the token with the largest logit, the model's score for that possible next token. If the calculation changes the scores, a close contest can have a different winner. Imagine serving the same prompt twice. The first time it is the only request waiting; the second time the server processes it beside several longer prompts. Padding, tile sizes, and the amount of parallel work may differ. We still want the selected request to receive the same answer. Fixing the prompt and the weights does not fix all those implementation choices.

Horace He and the Thinking Machines team state the requirement for their root-mean-square normalization (RMSNorm) reduction precisely:

the reduction order for each element must be fixed regardless of the batch-size of the kernel.

That gives us something specific to formalize. We can let the server group requests differently while requiring each selected row to keep its own addition tree. If the tree changes, we need a different argument, using an error bound and the winning margin.

A tile is a group of coordinates the kernel processes together. A different tile size can change where partial sums are combined, even when the selected row contains the same numbers. That is the implementation choice we'll represent explicitly below.

We have two questions to answer. We can keep the arithmetic fixed and ask for exactly the same scores. Or we can allow small numerical differences and ask whether the winning token stays the same. I'll work through both, then look at what happens when a server drafts several tokens and has to reject some of them.

Why moving the parentheses changes the sum

The smallest useful example has three numbers: a large positive number, its negative, and one. Let a be the binary32 approximation of 10²⁰. There are two ways to parenthesize their sum:

On the left, the large numbers cancel exactly, leaving one. On the right, the first addition tries to add one to a number whose magnitude is about 10²⁰. Binary32 cannot represent that small change at that scale, so rounding loses the one. The final addition gives zero.

Binary32 stores each number in 32 bits. As described in Goldberg's account of floating-point arithmetic, a normal binary32 number has 24 bits of precision, including the implicit leading bit. The exponent lets it represent large magnitudes, but does not give it extra significant digits. At a magnitude near 10²⁰, neighboring representable values are much farther apart than one, so adding one can leave the stored value unchanged.

All three inputs are identical in both calculations. Each addition rounds before the next one starts, and that makes the parentheses part of the program. Replacing a floating-point sum with an equal sum over the real numbers would erase the behavior I want to study.

In Core.lean, IEEE32.add_nonassoc_witness checks this inequality using FloatLib's executable binary32 model. The large operand has bit pattern 0x60AD78EC; its negative and one are specified by bits too. Lean can decide the concrete inequality without assuming that addition is associative.

Here are those two calculations in Lean. Hover over a name to see its type; the output below comes from running the commands, using the same operands as the theorem.

import BatchInvariantInference.Core

open BatchInvariantInference.IEEE32
open FloatLib.Floats

-- Use the same binary32 bit patterns as the non-associativity theorem.
#eval ExecFloat.add (ExecFloat.add big negBig) one
#eval ExecFloat.add big (ExecFloat.add negBig one)
#check BatchInvariantInference.IEEE32.add_nonassoc_witness

1

0

BatchInvariantInference.IEEE32.add_nonassoc_witness :
  (ExecFloat.add big negBig).add one ≠ ExecFloat.add big (ExecFloat.add negBig one)

The bit patterns matter because a decimal expression such as 1e20 is already an approximation when stored in binary32. Specifying the bits fixes the operand before we ask Lean to evaluate either tree. The model also fixes the rounded addition operation. The inequality is about those particular floating-point values and operations, rather than an unexplained difference between two decimal printouts.

Try the smaller magnitudes in the figure. At 1,024, the spacing of representable values is small enough for the one to survive. The 2²⁴ option is more subtle. Binary32 spacing changes at powers of two; the step from −2²⁴ toward zero is still one, so −a + 1 remains representable here. Changing the order of additions can matter, but it does not force a different answer for every choice of inputs. The Lean witness gives one definite failure.

Choosing the order of a sum

A dot product multiplies corresponding entries, then adds the products. A matrix multiplication performs many such reductions. A backend may split a long reduction across workers, combine partial sums, or choose a different kernel when the batch gets larger. Thinking Machines Lab traces batch-dependent numerics through matrix multiplication, RMSNorm, and attention.

I'll use “reduction” for combining several values into one, eg adding the products in a dot product. I represent its addition schedule as a tree. A leaf names an input; an internal node adds the results of its two children. The expressions above become two trees with the same three leaves. Choosing the tree is now an explicit input to the computation.

For four inputs, a left fold computes \(((x_0+x_1)+x_2)+x_3\). A balanced tree computes \((x_0+x_1)+(x_2+x_3)\). Both read each input once, in the same left-to-right leaf order, but the balanced version can compute the first two additions independently. A GPU can use that parallel work; it also rounds at different intermediate values.

ReductionTree comes from FloatLib. Its evaluator takes both a tree and a function for combining two values. That function could be rounded addition, exact rational addition, or another operation. The basic batch-invariance proof does not need associativity. If the selected leaves and the whole tree agree, every operation receives the same arguments. The proof can therefore keep the numerical schedule instead of replacing it with an unordered mathematical sum. Thats why I want the schedule to be an input too.

The counterexample gives a singleton batch the left tree and a larger batch the right tree. It keeps the selected row fixed, yet changes that row's result from one to zero. This small constructed schedule proves that batch-dependent arithmetic can violate the serving property; it does not identify the schedule used by a particular production kernel.

The positive result uses the same definitions. Equal selected rows evaluated with equal trees give equal results. A scheduling policy may depend on the request's hidden dimension, sequence length, dtype, or block size. The condition is that an unrelated request cannot change the selected row's schedule. The matmul and RMSNorm results spell out this condition for their respective reductions.

Two requests with different context lengths may use different trees. The condition only compares requests with the same local metadata. What it excludes is a tree change caused solely by someone else's request entering the batch. It does not prove that a particular server obeys this condition; the server's policy has to be connected to the definition.

Attention adds a second complication. The server can divide a sequence into chunks or put some keys and values in a cache while computing the rest. If those choices change where partial sums meet, they can change rounding too. The attention specification therefore includes the ordered key/value reduction tree. Its canonical-layout condition requires that the same request-local context give the same tree across chunk plans. The fixed-block construction satisfies that condition by deriving the block order from the request itself.

For example, consider a history of 1,000 keys. Dividing it among four workers gives four groups of 250; dividing it among eight gives groups of 125. That changes the intermediate sums. A fixed split size of 256 instead gives groups of 256, 256, 256, and 232, regardless of how many workers happen to be available. Workers can process those groups in parallel while retaining the specified order for combining their results. This is the fixed-size Split-KV strategy described by Thinking Machines Lab.

The reductions inside a transformer

Consider one output coordinate of a matrix product. Its leaves are the products of the selected input row and one weight column. Another batch row has no place in that formula. If the input row, weights, multiplications, and addition tree stay fixed, that output coordinate stays fixed. Repeating the argument for every column gives the matrix result. The tensor theorem represents the rows and weights with shapes checked in their types.

RMSNorm, introduced by Zhang and Sennrich, uses a reduction too. For a row of width d, it computes the mean of the squared coordinates, adds a small positive constant, and rescales the row:

Here wⱼ is a learned scale and δ is the normalization constant. The sum in the denominator is shared by every output coordinate. If its rounding changes, the change can affect the whole row. The theorem keeps the supplied reduction and scalar operations fixed, so equal rows produce equal denominators and equal normalized outputs. A runtime proof would also have to identify the square root and other scalar operations it implements.

Attention combines several reductions: query–key dot products, the normalization of attention weights, and weighted sums of value vectors. Preserving only the final value-sum tree would leave the earlier scores free to change. The attention definitions therefore state the feature schedules and ordered key/value layout used by the calculation. The proof follows those inputs through scores, weights, and outputs.

When score changes preserve the next token

A changed score can leave the next token unchanged. Suppose the best two scores are 8 and 3. Moving each score by at most 0.01 leaves a large gap. With scores of 8 and 7.999, the same error budget allows the order to reverse.

Let the winning margin be m, measured against every competing token, and bound the absolute change of each logit by ε. The winner can move down by ε while a competitor moves up by ε. The gap can shrink by at most 2ε:

winner_stable_of_margin proves this argument. choose_stable_of_margin connects it to a decoder that implements strict argmax. These results let me state a token-level guarantee without requiring the two logit vectors to be bit-for-bit equal. Applying the theorem to a runtime still requires a bound on that runtime's logit error.

We can see the proof without knowing any tactics. Write the reference winning score as zᵦ and one competitor as zₜ. In the perturbed calculation, the winner is at least zᵦ − ε, and the competitor is at most zₜ + ε. Subtracting gives:

The same inequality must hold for every competitor, not only a pair chosen after looking at the result. A uniform bound on all vocabulary coordinates supplies that premise. At equality, m = 2ε, the allowed intervals can touch. Then the theorem cannot guarantee a strict winner, and the decoder's tie-breaking rule may matter. This is why the hypothesis uses a strict inequality.

There is another error convention to watch. Suppose two implementations are each within ε of one shared reference. The triangle inequality bounds their distance from one another by 2ε. Applying the direct-perturbation theorem to that bound requires a margin greater than 4ε in the first implementation. Alternatively, if the reference itself has margin greater than 2ε, we can apply the theorem to each implementation separately and show both choose the reference winner. The code provides results for these different premises; calling them all “error ε” would hide which comparison we had proved.

Here we're comparing two executions of the same request, with the same input and weights. Classifier robustness asks a related question: does the winner survive changes to the input itself? The margin argument can help with either, once we have the appropriate logit bound. What supplies that bound is different. A numerical bound must follow rounding and execution choices; a robustness bound must follow the allowed input perturbation through the model.

I also wanted to see the size of these differences in a model, so we ran the FP16 probe above. The logits moved, but by much less than the gap between the two highest scores, and the winning tokens agreed. To make the same argument for later tokens, we need the margin and error bound at each new prefix: every selected token becomes part of the next input.

Each prompt number refers to one measured pair of executions. An observed maximum difference is a bound for that pair's recorded logits. It does not become a bound for all future requests or every kernel the server might select. I keep it beside the theorem because it makes the size of the margin concrete, while leaving the theorem's premise visible.

To carry the result through generation, start with equal prefixes. Equal prefixes give us the same reference input for the next step. If the step's margin and error satisfy the inequality, both implementations append the same token. We then have equal prefixes again and can repeat the argument. One failed margin condition interrupts this proof, even when the two programs happen to keep agreeing in a test.

From layer errors to logit errors

Now we need an error bound for the logits. It has to be small enough to leave every competing score below the winner, so let's follow an input error through a linear layer.

For a linear layer with fixed real weights, an input-coordinate error contributes through each weight that reads it. If \(y_j=\sum_i x_i W_{ij}\) and the second calculation uses the same weights, the triangle inequality gives

A numerical implementation adds local rounding error too. Let each input-coordinate error be at most \(E\), and let \(L\) bound every column sum \(\sum_i |W_{ij}|\). With a bound \(\delta\) for the local arithmetic, the next allowance is \(LE+\delta\): amplify the incoming error by at most \(L\), then add the local arithmetic allowance \(\delta\). A large weight can amplify an earlier difference; a small local rounding error alone doesn't control the final logits. Errors in stored weights would contribute additional terms, so the fixed-weight premise is doing real work here.

For several layers, the later maps can amplify errors introduced earlier. Their numerical contracts have to account for the actual operations and ranges, including normalization denominators. Once the resulting bound covers every logit at the current prefix, the margin theorem tells us whether it is sufficient for the next token. These equations explain the analysis required; this article doesn't construct a complete numerical certificate for a production transformer.

For a new backend, we can reuse the argmax proof once we establish its numerical bound. For a new scheduler, we can reuse the reduction-tree proof if it preserves the selected request's tree. In both cases, the work is checking the implementation against the hypotheses of the theorem we already have.

Committing a draft without corrupting the cache

A serving system can speculate ahead, verify a window of proposed tokens, and discard rejected candidates. Leviathan, Kalman, and Matias describe speculative decoding and prove a distribution-preserving sampling procedure. Here I wanted to check the committed state and tokens, since a user receives the tokens that survive verification and are committed.

The invariant relates two pieces of state: the committed token list and the reference decoder state after those tokens. Speculation changes only the fast-path state and candidates. Rollback clears candidates. A commit must carry evidence that the accepted window and resulting reference state match an execution of the reference decoder. An induction over these transitions proves ServeDVR_refines_RefDecode: every committed output is a reference continuation prefix. Two legal traces that commit the same number of tokens therefore show the same tokens to the user.

For example, a draft might propose three tokens while the verifier accepts only the first two. The committed prefix advances by those two tokens; the rejected third token must not advance the reference state. A later draft starts from the reference state after the accepted pair. An implementation that discards the token but keeps its cache update would violate the invariant, because its state would describe a longer history than the user has received.

The proof does not need to know how clever the draft is. Speculation may change its private state and produce poor candidates. What matters is the commit rule: the verifier supplies agreement for the accepted window and its resulting reference state. Rollback leaves that committed state alone. Induction over a serving trace then applies the same local argument after every speculation, rollback, or commit.

This also explains the “same number of tokens” premise. Two servers may stop at different times or commit different-length windows. Their outputs can be different-length prefixes of the same reference continuation. The theorem gives equality once those lengths agree; it does not prove equal latency or an equal number of completed tokens at a given instant.

The verifier's agreement with the reference is a hypothesis of this theorem. The development also connects the serving statement to our IR decoder, whose reference computation is the denotation of a model graph. There is a smaller experiment closer to compiled code as well. The CUDA reduction checker accepts a certificate extracted from an eight-step fused multiply-add chain and proves that the accepted chain denotes the intended value reduction. Applying it to a live kernel requires the stated indexing, allocation, and launch conditions. The certificate covers that reduction fragment; a full attention kernel needs the remaining computation connected to the attention specification.

A fused multiply-add rounds the combined product and addition once. Separate multiplication and addition round twice. The CUDA fragment therefore has its own FMA-chain semantics; we should not silently identify it with every reduction written using separate operators. The accepted certificate connects the extracted chain to that specified reduction. Its indexing and launch hypotheses are part of the statement, since arithmetic agreement alone cannot show that a kernel loaded the intended coordinates.

For a larger serving implementation, I would first identify the request-local reduction policies, then establish numerical or exact agreement for the layer computations, and only then apply the token and commit theorems. The small certificate makes one of those steps inspectable. The remaining layer and runtime connections still have to be supplied.

Conclusion

I'd like to take this further with an inference implementation. We would need its actual reduction schedules and cache policy, then a bound on the score differences at each reached prefix. The margin theorem tells us how to use that bound; obtaining it for the running model is the next job. If you're working on those kernels or serving policies, I'd love to compare examples :)

References

• Horace He and collaborators at Thinking Machines Lab. Defeating Nondeterminism in LLM Inference, 2025. The batch-invariance problem and fixed-size reduction strategies that motivated this example.
• David Goldberg. What Every Computer Scientist Should Know About Floating-Point Arithmetic, 1991. Representation, rounding, and non-associative arithmetic.
• Biao Zhang and Rico Sennrich. Root Mean Square Layer Normalization, 2019.
• Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast Inference from Transformers via Speculative Decoding, 2023. The original sampling result is distinct from our state-and-commit theorem.

@misc{george2026batch,
  author = {Robert Joseph George},
  title = {{Thinking Machines Lab’s batch-invariance problem, formalized}},
  year = {2026},
  howpublished = {TorchLean Verified Examples},
  url = {https://robertboy18.github.io/TorchLean-Verified-Examples/week-01-batch-invariant-inference/}
}