Deterministic Inference across Tensor Parallel Sizes That Eliminates Training-Inference Mismatch

Deterministic inference is increasingly critical for Large language model (LLM) applications such as LLM-as-a-judge evaluation, multi-agent systems, and reinforcement learning (RL). However, existing LLM serving frameworks can produce different outputs for identical inputs when tensor parallel (TP) size or batch size changes, even under greedy decoding. This arises from the non-associativity of floating-point arithmetic and inconsistent reduction orders across GPUs. While prior work has addressed batch-size–related nondeterminism through batch-invariant kernels, determinism across different TP sizes remains an open problem, particularly in RL settings, where the training engine typically uses Fully Sharded Data Parallel (FSDP) (i.e., TP = 1) while the rollout engine relies on multi-GPU TP to maximize the inference throughput, creating a probability mismatch that can degrade or even collapse training. We identify and analyze the root causes of TP-induced inconsistency and propose Tree-Based Invariant Kernels (TBIK) , a set of custom matrix multiplication and reduction kernels that guarantee bit-wise identical results across TP sizes. Our key insight is to enforce a consistent reduction order across and within GPUs. We implement TBIK in Triton and integrate it into vLLM and FSDP, achieving bit-wise deterministic inference across different TP sizes and zero probability divergence between rollout and training engines in RL pipelines. By eliminating mismatches caused by different parallelization strategies, TBIK enables true on-policy RL at scale for the first time , leading to improved model performance and faster convergence.