Bit-Exact After the TP Shape Changes?

Created on September 20, 2026 by Yanying Lin

2026   ·   Tensor Parallelism   Deterministic Inference   vLLM   NCCL   ·   Research   Systems

TP1 has no collective, so bit-exact agreement across tensor-parallel shapes can only come from one shared tree for the on-GPU K reduction and the cross-GPU reduction. TBIK already builds that tree.

Where deterministic inference stands

Floating-point addition is not associative. In bf16, (a+b)+c and a+(b+c) often differ in the last bit. Every matmul and every softmax in a large-model forward pass sums thousands of terms. Change the order and the last bit moves. Under greedy decoding, a top-1 that flips between two close candidates sends the rest of the continuation down another path. For a long time this was treated as unavoidable noise.

In 2025, Thinking Machines pointed out that the main source of nondeterminism in inference serving is not a race among GPU threads. It is the batch size changing the kernel’s tile choice and split, and therefore the summation order. They published a set of batch-invariant kernels with a fixed tile, so the same request yields the same bits no matter which other requests share its batch (He et al., 2025). vLLM then added VLLM_BATCH_INVARIANT. The switch replaces matmul, normalization, and related operators, pins NCCL to the tree algorithm, one channel, and the Simple protocol, and turns off the custom allreduce. It requires compute capability 8.0 or newer (vLLM batch invariance). SGLang’s deterministic mode takes the same path. It fixes custom allreduce to a single-stage kernel that sums in rank order, and on H100 it is about 34% slower on average (LMSYS, 2025).

These modes provide several invariants that are independent of each other: the same bits from run to run at a fixed shape, the same bits across batch sizes, and sometimes the same bits between prefill and decode. A different GPU model and a different parallel shape are two further axes. Getting the first few does not hand you the last two.

Why the parallel shape changes

Tensor parallelism (TP) splits one layer’s weights across GPUs. The Megatron split cuts the first MLP matrix by columns. Each GPU computes some of the output columns and needs no communication. The second matrix is cut by rows, along the reduction dimension K. Each GPU produces a partial sum with the full output shape but only a slice of K, and an allreduce adds those partial sums (Shoeybi et al., 2019). The attention output projection works the same way. An 8B dense model with hidden size 4096 and intermediate size 14336 has two row-parallel reductions in every layer.

The TP degree changes in real systems. A serving stack scales instances with load, and a request can move from a TP4 instance to a TP1 instance. In reinforcement learning, the rollout side and the trainer side often use different parallel configurations. The probability of the same trajectory is computed twice, and any numerical gap becomes off-policy bias. The usual response is an importance correction, or bit-exact alignment at one fixed shape. The natural question is whether the deterministic promise still holds when the TP shape changes.

TL;DR

Deterministic inference already makes one fixed deployment shape repeat bit for bit. The natural guess is to extend that to a changing tensor-parallel degree, so TP1, TP2, and TP4 emit the same bits.

The hard part is not the collective. A row-parallel layer splits K by the TP degree and rounds each piece. TP1 has no collective, so tuning NCCL can never make TP1 agree with TP4. The on-GPU K reduction and the cross-GPU reduction have to share one fixed tree.

TBIK (arXiv 2511.17826) already builds that tree on vLLM and reports bit-exact TP 1/2/4/8, at 22% to 63% end to end against native BF16, and only for power-of-two TP and row-parallel layers. The tree is already there. What remains is extension and a lower cost. It is still unclear who actually needs cross-shape agreement. The conjecture fails before any GPU measurement.

The conjecture we can test

The starting judgment was that a numerical result only has to be stable inside one deployment shape, and that mathematical equivalence is enough across shapes. Elastic serving and separated training and inference make that old assumption too weak. The plan was a canonical arithmetic DAG that does not depend on the TP shape. Split the K dimension of each row-parallel layer into a globally fixed number of chunks C, where C is a common multiple of the TP degrees we support. Reduce each chunk on a fixed tree. TP1 simulates the same split and the same rounding points inside one kernel. When TP is larger than 1, each GPU computes its own chunks and the GPUs merge them on the same tree. This looked feasible because a fixed-order exact mode on one GPU is already cheap, and pushing it to many GPUs looked like one extra convention on the collective.

The split is on the GPU, not in the collective

The existing switches are not enough, and the reason is visible before any measurement. Figure 1 is the same sum of K=8 tiles under three executions.

Three reduction shapes for the same eight K tiles: TP4 local sums plus a ring, a single linear K loop on TP1, and one shared tree used by both.
Figure 1. Three reduction shapes for the same K sum. A sketch, not a measurement. (a) Default TP4 sums and rounds on each GPU, then a ring allreduce whose order depends on the number of ranks. (b) Default TP1 walks K in one kernel and has no collective. (c) A shared tree: the same binary tree on one GPU and across GPUs. On TP1 the dashed levels run inside that one GPU.

(a) and (b) differ in three places: where the split falls, how many intermediate roundings occur, and the order of the adds. Pinning NCCL’s algorithm, protocol, and channel count can make (a) repeatable for one rank count and one topology. It cannot change (b), because TP1 has no collective to configure. Any plan that only touches the communication library therefore cannot make TP1 agree with TP4. The GEMM kernel on the TP1 side has to adopt the same K split and the same tree as TP4. That is (c). This follows from the reduction structure. It does not depend on a measurement.

We estimated the cost of (c) with a closed-form model (Figure 2). The model uses the Llama-3-8B shape and nominal A100 PCIe parameters: 1.5 TB/s of HBM, 312 TFLOPS of bf16, 25 GB/s of effective cross-GPU PCIe, 15 microseconds of latency per collective, and C=4.

Analytic cost of a canonical K split. Left: extra TP1 GEMM time if four partial sums are written out. Right: TP4 all-gather plus a fixed fold, relative to a ring allreduce.
Figure 2. Cost of a canonical K split. A closed-form estimate, not a measurement. Left: extra time in the row-parallel GEMM when TP1 writes four partial sums and reads them back. Right: communication time of an all-gather plus a fixed-order local fold on TP4, relative to a ring allreduce.

The left panel says the TP1 cost grows with the batch. Decode at batch 1 is negligible, about 0.1%. Batch 64 is about 5%. An 8K-token prefill is about 18%. If the rounding at the split boundaries is fused into one kernel, there is no extra memory traffic. The right panel says the real cost is multi-GPU communication. An all-gather delivers every partial sum to every GPU, so the volume is about the rank count times an allreduce. For small messages the latency term dominates, and this path is faster than the two rounds of a ring: about 0.51× at decode batch 1 and 0.89× at batch 64. For large messages the bandwidth term dominates, and an 8K prefill is about twice the ring. The estimate assumes PCIe and nominal bandwidth. The ratios move on an NVLink machine.

Prior work already built this tree

In the search we ruled out the existing systems one by one. The deterministic modes in vLLM and SGLang, a pinned NCCL configuration, Megatron’s deterministic mode, and the true on-policy alignment in RL frameworks all hold only inside their own fixed parallel shape. Neither the docs nor the code claim agreement across TP shapes. ReproBLAS uses binned floating point so a sum does not depend on order, which removes the problem at the root, but it is CPU BLAS and MPI reductions (Demmel and Nguyen, 2015).

Searching by mechanism turned up Tree-Based Invariant Kernels, TBIK (arXiv 2511.17826). Its TP-invariant matmul is a persistent Triton kernel. It reduces K by tiles on a binary tree, and the upper levels of that tree are exactly the cross-GPU cuts at different TP degrees. The tree TP1 computes inside one GPU is the same tree TP4 assembles from on-GPU pieces plus the cross-GPU levels. That is Figure 1(c). The cross-GPU part is a tree allreduce. Small tensors use a custom kernel over IPC. Large tensors fall back to an all-gather and then add pairs in a fixed hierarchy. The paper explains why it does not use NCCL: inside a node, NCCL is still a chain reduction. The implementation is a patch that replaces RowParallelLinear in vLLM 0.11, switched on by two environment variables. It reports bit-exact agreement, in BF16, for Qwen3, Llama-3.1, Mistral, and related models, at TP 1/2/4/8 and batch sizes 8/16/32. The same kernels are used on the FSDP training side, so TP1 training and multi-GPU rollout probabilities differ by zero. The end-to-end cost against native BF16 is 22% to 63%. The tree matmul alone is 2% to 25%, and the tree allreduce is at most 10%. The boundary is stated clearly. TP must be a power of two. Only row-parallel layers are covered. Attention requires chunked prefill to be off. There are no results for MoE, pipeline parallelism, or fp16.

Why the conjecture does not hold

For the conjecture to stand, existing systems would have to be unable, in principle, to agree across shapes. TBIK refutes that, and it refutes it with the mechanism we were about to propose: one tree covering both levels of the reduction. What is left is TP degrees that are not powers of two (a fixed chunk count plus a linear fold does not require a power of two), expert parallelism for MoE, and a lower overhead. Those are extensions and engineering on the same mechanism. They are not a new capability.

The second reason is demand. The migration cases we wanted to serve are single GPU to single GPU, and a single-GPU exact mode already covers them. Aligning training and rollout is the use TBIK itself chose, and same-shape alignment plus an importance correction is already common practice. We could not find a user for whom bit-exact agreement across TP shapes is irreplaceable.

One gap remains. TBIK’s consistency claim comes from its paper and its code. We did not reproduce it on our own hardware. The Volta GPUs we have do not meet the compute-capability requirement of vLLM’s deterministic mode, and they have no native bf16. The only way that gap would hurt the technical conclusion is if TBIK failed under a wider set of conditions, and that failure sits in the extensions already listed above.

What this analysis supports directly

Bit-exact agreement across TP shapes, on a row-parallel layer, requires the on-GPU K reduction and the cross-GPU reduction to be the same logical DAG. TP1 has no collective. Pinning NCCL’s algorithm, protocol, and channel count does nothing. To judge whether a deterministic scheme can cross shapes, look first at whether it specifies the tree inside the K dimension.

Determinism has at least five independent axes: run to run, batch, prefill versus decode, GPU model, and parallel shape. Having one does not give you another. When you read the documentation of a deterministic mode, check which axis it actually claims.

The numbers below are a closed-form estimate, on PCIe at nominal bandwidth. The main cost of cross-shape agreement is not on the TP1 side. Once the rounding is fused into one kernel, a single GPU barely adds memory traffic. The cost is on the multi-GPU side, where an all-gather replaces an allreduce. Communication for a large-batch prefill roughly doubles. Small-batch decode is faster. The crossover is set by the message size and the latency of one collective.

Among existing open implementations, the deterministic modes in vLLM and SGLang give same-shape agreement only. TBIK gives cross-shape agreement for power-of-two TP, at 22% to 63% end to end against native BF16. That denominator is not a same-shape deterministic mode. A comparison has to use the same one.

Searching for deterministic inference by the name of the problem finds systems that share the name and not the invariant. Searching by the mechanism finds the work that actually covers it. Once the mechanism can be written in one sentence, search for that sentence before any characterization experiment.

TBIK already covers the defining case

What stays is the search result. TBIK already builds this tree on vLLM and reports bit-exact TP 1/2/4/8. Against native BF16 the end-to-end cost is 22% to 63%, and only for power-of-two TP and row-parallel layers. The tree itself is already in prior work. What remains is extending it to more shapes and lowering the cost. It is still unclear who cannot do without bit-exact agreement across shapes. The conjecture fails before any GPU measurement.