Batch-invariant inference kernels (Thinking Machines)

Research demonstrationProvider-reported

Batch-invariant kernels are GPU routines that make a language model give the same output for the same input, whatever else the server is processing.

Thinking Machines Lab identified changing batch sizes as the main reason language-model services are nondeterministic, and published kernels that combine partial results in an order independent of the batch. With them, 1,000 temperature-zero completions from a 235-billion-parameter model were identical. The inference engine SGLang built its deterministic mode on the kernels, and vLLM's batch-invariant mode, still in beta, is based on the same work.

Identical outputs let a verifier that runs the same model, engine and hardware require an exact match instead of a tolerance. The kernels were built for reproducibility, not verification, and no verification system is documented using them.

They made inference take 61.5% longer in Thinking Machines' Qwen3-8B test, whose vLLM integration was not yet heavily optimised. SGLang's tests averaged 34% longer latency.

Development assessmentLow confidence

Public kernels made a 235-billion-parameter model's outputs identical across 1,000 runs and underpin the deterministic modes of two production engines, but no verification system is documented using them.

Assessed use: exact recomputation of served outputs by a verifier, with a cooperating provider

Rubric assessment and gaps →

On this page

What it is

Thinking Machines Lab published batch-invariant kernels for large language model (LLM) inference in September 2025 1. Its post argues that "the primary reason nearly all LLM inference endpoints are nondeterministic is that the load (and thus batch-size) nondeterministically varies" 1. The companion library, batch_invariant_ops, is MIT-licensed 2.

The kernels are one route to deterministic and bit-exact inference. The SGLang inference engine built its deterministic mode on them 3, and the vLLM engine's developers state in a GitHub tracking issue that its batch-invariant mode is based on the post 5.

How it works

A kernel is batch-invariant if it computes each output element with the same reduction order whatever the batch size 1. The post identifies three operations that involve reductions: RMSNorm, matrix multiplication and attention 1. For attention it fixes the size of each split along the key-value dimension, rather than the number of splits 1.

The library replaces PyTorch's matrix multiplication (torch.mm and torch.addmm), log-softmax and mean through torch.Library, so existing models can use the kernels with small code changes 2. The post runs the result in vLLM through its FlexAttention backend 1. vLLM now exposes its own mode with the environment variable VLLM_BATCH_INVARIANT=1, on NVIDIA GPUs of compute capability 8.0 or higher, AMD Instinct GPUs and Intel XPUs 4.

For verification, identical outputs hold when the model, inference implementation and device are fixed 6. A verifier that runs the same stack can then require an exact match, instead of the tolerance that sampled recomputation schemes allow.

Search

Full search page