Thinking Machines Lab - Defeating Nondeterminism in LLM Inference: The Core Role of Batch Invariance

Defeating Nondeterminism in LLM Inference - Thinking Machines Lab

EN

Unveiling Nondeterminism in LLM Inference: The Path to Reproducible AI

In the field of artificial intelligence, LLMs have become indispensable tools. However, a persistent challenge for researchers and application developers is that despite the theoretical determinism of underlying computational hardware, LLM inference results often exhibit nondeterminism. This means that even with identical inputs and a temperature setting of 0 (greedy sampling, which should theoretically produce deterministic results), each run may yield different outputs. This lack of reproducibility, as pointed out by Horace He and Thinking Machines Lab, is a cornerstone of scientific progress but remains exceptionally difficult to achieve in the LLM domain, significantly hindering its development in scientific research, particularly in applications with strict determinism requirements such as "True on-policy RL".

To intuitively understand this phenomenon, we can observe that even with the sampling temperature reduced to 0, LLM APIs may still produce nondeterministic results in practice. For instance, the probability distribution generated by the DeepSeek V3.1 model at temperature=0 for the same prompt shows that even the highest-probability word may not be 100% certain, and other suboptimal options might also be selected, leading to diversity in the final output.

LLM Generated Probability Distribution Example

Source: Thinking Machines Lab

Floating-Point Non-Associativity: Root Cause of Numerical Differences, Not Direct Cause of Nondeterminism

For a long time, the industry generally believed that nondeterminism in LLM inference stemmed from a combination of floating-point non-associativity and concurrent execution. Floating-point non-associativity refers to the fact that due to limitations in computer floating-point representation, addition operations do not satisfy the associative property, meaning .

For example, in floating-point arithmetic,  might result in 0, while  yields 0.1. This occurs because floating-point numbers lose precision when adding numbers of vastly different magnitudes. For instance, when adding 1230 and 23.4, if only 3 digits of precision can be maintained, the result will be 1250, not the precise 1253.4. This implies that each time floating-point numbers are added in a different order, a different result may be obtained.

Precision Loss in Floating-Point Addition

Source: Thinking Machines Lab

However, experiments by Thinking Machines Lab demonstrate that even when repeating the same matrix multiplication on a GPU, the results are consistently bit-for-bit identical. This suggests that while floating-point non-associativity is the fundamental cause of numerical differences, it does not directly explain the true source of LLM inference nondeterminism, nor does it clarify how or when floating-point numbers are added in different orders, or how to avoid it. The answer lies deeper, in the implementation of kernels and the working mechanisms of the inference system.

Batch Invariance: The True Culprit Behind LLM Inference Nondeterminism

Although GPU kernels are generally "run-to-run deterministic"—meaning they will always provide the same deterministic output given identical inputs—this alone is insufficient to ensure that the entire LLM inference system containing them is also deterministic.

From the perspective of an inference server, the model itself is deterministic when receiving a specific user request; its forward propagation will be consistent.

Model Determinism vs. User Request

Source: Thinking Machines Lab

However, when multiple user requests arrive concurrently at an inference server, they are typically batched together for processing to improve GPU utilization and throughput. The size of this batch dynamically changes based on server load. The true culprit behind LLM inference nondeterminism is the lack of "batch invariance". This implies that when the batch size changes, the computational result for each individual request within the batch may vary.

Model Determinism vs. Nondeterministic User Request

Source: Thinking Machines Lab

For example, in matrix multiplication, calculating the first element of a batch individually may yield a significantly different result compared to calculating the entire batch and then extracting the first element. This is because batching alters the data access patterns and reduction order on the GPU. Since user requests arrive randomly, the load experienced by the inference server is nondeterministic, which directly leads to nondeterministic changes in batch size, consequently altering the final result of each independent request. This nondeterminism is not exclusive to GPUs; it is also prevalent in LLM inference endpoints on CPUs or TPUs. Therefore, the primary reason for nondeterminism in almost all LLM inference endpoints is the nondeterministic variation in load (and thus batch size).

Achieving Batch-Invariant Kernels: A Core Strategy to Overcome Nondeterminism

To achieve batch invariance for Transformer models, it is crucial to ensure that every core operation within the model possesses batch invariance. Pointwise operations are typically batch-invariant, but three core operations involving reduction—RMSNorm, Matrix Multiplication, and Attention—require special handling.

1. Batch-Invariant RMSNorm

RMSNorm is a common normalization method whose computation involves root mean square normalization of the input tensor. The key to achieving batch-invariant RMSNorm is that the reduction order for each element must remain fixed, regardless of changes in batch size.

The standard data-parallel strategy assigns each batch element to a GPU core, ensuring that each reduction is completed entirely within a single core.

Data-Parallel RMSNorm

Source: Thinking Machines Lab

However, when the batch size decreases, there may be more cores than batch elements. To maintain parallelism and performance, engineers sometimes employ a "split-reduction" strategy, where the reduction task for a single element is distributed across multiple cores. This, however, changes the reduction strategy and order, thereby breaking batch invariance.

Split-Reduction RMSNorm

Source: Thinking Machines Lab

The typical solution is: for small batch sizes, sub-optimal performance can be accepted as kernel execution is usually fast. Alternatively, consistently use a reduction strategy that provides sufficient parallelism even for very small batches but maintains a fixed reduction order, even if this means sub-optimal performance for larger batches.

2. Batch-Invariant Matrix Multiplication

Matrix multiplication is a core operation in LLMs. Its standard parallel strategy is also data-parallel, where the output is tiled into 2D tiles, and each tile is assigned to a different core.

Data-Parallel Matrix Multiplication

Source: Thinking Machines Lab

However, when batch dimensions (M and N) are too small, leading to insufficient parallelism, it may be necessary to split along the reduction dimension (K), using "Split-K Matmul". This introduces a new reduction order, thereby breaking batch invariance.

Split-K Matmul

Source: Thinking Machines Lab

Furthermore, the use of Tensor Core instructions adds complexity. Different Tensor Core instructions (e.g., for different tile sizes) may lead to different internal reduction orders. When batch sizes are very small, it might be necessary to switch to smaller Tensor Core instructions or abandon Tensor Cores entirely, both of which break batch invariance.

To ensure batch invariance for Matrix Multiplication, the simplest approach is to compile a single kernel configuration and use it for all shapes, even if this sacrifices some performance. In LLM inference, since the model dimension N is typically large, this performance loss is usually not catastrophic.

Batch-Invariant Matrix Multiplication Performance

Source: Thinking Machines Lab

The chart above compares the performance of the batch-invariant kernel (orange line) with cuBLAS (blue line). Although the batch-invariant kernel incurs a performance loss of approximately 20%, and performance significantly drops at very small batch sizes, the trade-off of performance for determinism is worthwhile.

3. Batch-Invariant Attention

The Attention mechanism introduces two additional complexities: the reduction dimension includes not only the feature dimension but also the sequence dimension; and it must handle various inference optimizations that affect sequence processing, such as chunked prefill and prefix caching.

The data-parallel strategy introduced by FlashAttention2, which parallelizes along the query (Q) dimension while reducing keys (K) and values (V), is standard practice.

FlashAttention2 Strategy

Source: Thinking Machines Lab

However, the challenge lies in ensuring that the reduction order for a given token does not depend on how many other tokens are being processed concurrently in its sequence, such as the presence of a KV cache. If the K/V values in the KV cache are processed separately from the currently processed K/V values, batch invariance will be broken.

FlashAttention with KV Cache

Source: Thinking Machines Lab

Furthermore, in LLM inference, attention shapes often require split-reduction kernels (e.g., Split-KV or FlashDecoding) because the query length in the decoding phase is very small, leading to insufficient parallelism to saturate the GPU. This also breaks batch invariance.

Fixed Number Split KV Strategy

Source: Thinking Machines Lab

The ultimate solution is to adopt a "fixed split-size" strategy rather than a fixed number of splits. This means fixing the size of each split, which leads to a variable number of splits. This guarantees that the reduction order remains consistent regardless of how many tokens are processed, thereby achieving batch invariance.

Fixed Size Split KV Strategy

Source: Thinking Machines Lab

Experimental Validation and Future Prospects

Thinking Machines Lab successfully demonstrated deterministic inference by leveraging vLLM's FlexAttention backend and torch.Library. Their implemented batch-invariant kernel library is open-sourced in the thinking-machines-lab/batch-invariant-ops GitHub repository.

The experimental results are encouraging:

• Complete Elimination of Nondeterminism: Using the Qwen/Qwen3-235B-A22B-Instruct-2507 model, 1000 completions were generated for the same prompt at temperature=0. Results showed that before enabling batch-invariant kernels, 80 unique completions were generated, with the first divergence occurring at the 103rd token. After enabling batch-invariant kernels, all 1000 completions were identical, proving their effectiveness.
• Performance Trade-offs: While deterministic inference brings significant stability, there is currently some performance overhead.
  • Running Qwen-3-8B on an API server, completing 1000 sequences:
    • Default vLLM configuration took 26 seconds.
    • Unoptimized deterministic vLLM took 55 seconds.
    • With improved attention kernels, time was reduced to 42 seconds.
  • The performance degradation primarily stems from the fact that FlexAttention's integration into vLLM has not yet undergone extensive optimization.

Inference Performance Comparison

• Enabling True On-Policy Reinforcement Learning: Deterministic inference is crucial for reinforcement learning (RL). Numerical differences between training and inference can implicitly transform on-policy RL into off-policy RL, affecting training effectiveness. Through deterministic inference, bit-for-bit identical results can be obtained between the sampler and the trainer, enabling true on-policy RL.
  • Experiments in the RLVR setting using Bigmath and Qwen 2.5-VL instruct 8B showed that without importance weighting (off-policy correction), rewards would collapse midway through training. Adding off-policy correction terms or achieving true on-policy RL allowed training to proceed smoothly.
  • The KL-divergence plot shows that with importance weighting, KL-divergence remained around 0.001; without weighting, KL-divergence sharply increased around Step 318; whereas true on-policy RL kept KL-divergence at 0, indicating no divergence between policies. This opens new possibilities for RL research and applications.

Reinforcement Learning Reward vs. KL-Divergence

Source: Thinking Machines Lab

Conclusion

The research by Thinking Machines Lab strongly refutes the "defeatist" attitude often adopted when facing nondeterminism in machine learning. By deeply understanding and addressing the core issue of batch invariance, nondeterminism in LLM inference can be overcome. This breakthrough not only provides a solid foundation for reproducibility in LLMs for scientific research and AI applications but also encourages the entire community to more comprehensively understand and optimize their systems, thereby driving the further development of artificial intelligence technology.