Batch-Invariant Kernel Design: Conv2d and Attention BMM

View the code on GitHub

Introduction

In our previous post last week, we discussed how we made π0 inference batch invariant. Two operations required model-specific changes. SigLIP's patch projection used Conv2d, which was not batch invariant. Attention used batched matrix multiplication, which was not covered by the existing aten::mm replacement.

We implemented both operators directly and opened a PR in Thinking Machines' batch_invariant_ops.

Why does this matter? The math of a forward pass is run-to-run deterministic¹ and does not depend on which other inputs share its batch. GPU execution is not. The number of requests an inference server handles at once changes the batch size from run to run, and with it the reduction order inside each kernel, so the same input can come out with slightly different numbers depending on its batchmates.

In LLM inference, a small numeric difference can flip which token gets sampled. In a VLA, the same divergence introduced early in the model compounds through the layers and can select a different robot action. That is why VLA inference must be batch invariant.

We implement a per-sample Conv2d for SigLIP's patch projection and a Triton BMM for attention. In both kernels, the reduction for each output is independent of batch size. We verify this by showing that sample 0 produces bitwise-identical outputs at batch sizes 1 and 32.

The VLA trace exposed two missing kernels

In our previous blog post, we traced π0, a vision-language-action inference path, while comparing a reference example alone against the same example at index zero in a batch of size 32. We found two operations that were breaking batch invariance:

1. 01The SigLIP patch embedder nn.Conv2d: In vision-language models, the patch embedder cuts the input image into patches and turns each patch into a token vector. In π0/SigLIP models, the patch embedder usually uses a 14 × 14 convolution and performs a large reduction over image channels and patch elements. The problem is that, in PyTorch, the reduction strategy on the GPU is scheduled by the nn.Conv2d layer depending on the input batch size.
2. 02Joint-model attention torch.matmul: Joint-model attention flattens the batch and head dimensions and computes attention using rank-3 tensors. Although the attention module calls torch.matmul with Thinky’s batch-invariant mode on, PyTorch still dispatches aten::bmm instead of aten::mm by default, which breaks batch invariance.

We implemented and pushed both missing paths in the proposed batch_invariant_ops change². The implementation follows the same principle as the original persistent GEMM: make the reduction scope for a particular output independent of the input’s batch size.

Our kernel design ensures bitwise equality between the reference and batched input tensor:

Kernel 1: Conv2d

The first kernel is designed to perform convolution by iteratively unfolding each input image in the batch into its patch embeddings, and then multiplying by a flattened convolution filter matrix, using the batch-invariant GEMM. This solves the problem of having large reductions and the non-deterministic GPU reduction strategy that comes with it.

for sample in input:    patches = unfold(sample)             # [Cin × Kh × Kw, locations]    output  = flattened_weight @ patches # [Cout, locations]

Note that we loop through each sample in the input batch. This allows every sample to perform its own unfold and fixed-shape GEMM. Here’s a walk-through of our kernel design:

This is a generalization of the _UnfoldConv2d kernel in the π0 reference implementation, but it is registered for aten::convolution in Thinky’s batch-invariant ops library so it can be used by any model. Although it is not optimized for maximum throughput, it supports regular non-transposed 2-D convolution, including stride, padding, dilation, and groups, and is designed for stable reduction.

Kernel 2: Attention BMM

A typical attention score calculation has rank-3 shapes like this after the leading dimensions B, H are flattened:

Q:  [B × H, T, D]Kᵀ: [B × H, D, T]Q @ Kᵀ → [B × H, T, T]

The default aten::bmm that was dispatched earlier processes the entire input [B×H, T, D] @ [B×H, D, T]. Therefore, when the batch size changed, the backend either picked a new kernel or tile layout or changed the reduction strategy, which is not ideal for achieving deterministic outputs.

Our new BMM kernel solves this by launching a 2D grid. The X-axis represents the persistent tile workers, and the Y-axis represents the BMM elements, i.e., each element of the flattened B×H dimension. Each element gets its own base pointers and the same ordered reduction over K. This way, increasing B adds independent grids and does not modify the existing GEMM implementation.

The two correctness tests

Any new kernel that replaces PyTorch’s default ops must satisfy both numerical equivalence and the stronger batching property. To ensure these checks pass, we:

1. 01compare the replacement output against standard PyTorch within a tolerance
2. 02compare sample zero alone against sample zero in a larger batch and require an exact match in batch-invariant mode.

Although certain GPU + PyTorch + cuDNN versions might unconditionally assert that there is zero numeric divergence, this will not guarantee batch invariance for every run unless the kernels have a fixed reduction strategy. For that reason, we ran (and passed) both correctness tests on our kernel implementations.

{  "batch_size": 32,  "reference_samples": 1,  "distractor_samples": 31,  "num_trials": 1,  "reference_seconds": 1.2700269939377904,  "batch_seconds_mean": 23.618852018378675,  "batch_seconds_std": 0.0,  "throughput_speedup_vs_reference": 1.7206959836314306,  "metrics_mean": {    "max_abs_diff": 0.0,    "mean_abs_diff": 0.0,    "rmse": 0.0,    "relative_l2_mean": 0.0,    "relative_l2_max": 0.0,    "output_abs_max": 1.0  },  "metrics_worst": {    "max_abs_diff": 0.0,    "mean_abs_diff": 0.0,    "rmse": 0.0,    "relative_l2_mean": 0.0,    "relative_l2_max": 0.0,    "output_abs_max": 1.0  }}

The full trace has been logged and can be found in our previous blog post.

Discussion

It is important to note that our proposed kernels do not make an entire VLA inference bitwise deterministic. Softmax, norms, compiler fusion, cache operations, sampling, and every other reduction still need to be traced and checked individually.

Furthermore, since the Conv2d implementation trades some batching efficiency for iterating through inputs, we also do not guarantee that per-sample Conv2d is automatically the best performance choice.

The broader lesson here is that batch invariance must be treated as an operator-level property. A good practice is to start from a batch-difference trace, inspect the exact dispatched ATen operation and tensor ranks, then design a kernel whose accumulation order for one request cannot be changed by batch size.

Citations

Cybernetic Physics. “Batch-Invariant Kernel Design: Conv2d and Attention BMM.” Cybernetic Physics, September 2026.

@article{cyberneticphysics2026batchinvariantkernels,  author  = {Cybernetic Physics},  title   = {Batch-Invariant Kernel Design: Conv2d and Attention BMM},  journal = {Cybernetic Physics},  year    = {2026},  month   = {September}}

References

1. 01He, Horace, et al. “Defeating Nondeterminism in LLM Inference.” Thinking Machines Lab, September 2025. https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/
2. 02“Add batch-invariant Conv2d and BMM.” batch_invariant_ops, pull request 27. https://github.com/thinking-machines-lab/batch_invariant_ops/pull/27
3. 03Cybernetic Physics. “Batch-Invariant VLAs.” Cybernetic Physics, 2026. Read the previous blog post.
4. 04Thinking Machines Lab. batch_invariant_ops. GitHub, 2025. https://github.com/thinking-machines-lab/batch_invariant_ops
5. 05Black, Kevin, et al. “π0: A Vision-Language-Action Flow Model for General Robot Control.” Physical Intelligence, 2024. https://www.physicalintelligence.company/download/pi0.pdf
6. 06Tillet, Philippe, et al. “Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations.” MLSys, 2019. https://openai.com/index/triton/

Acknowledgements

We thank Thinking Machines Lab for publishing the original batch-invariant operator work and the persistent GEMM implementation in batch_invariant_ops. We also acknowledge Physical Intelligence for the π0 architecture, and the PyTorch and Triton communities whose tools made this implementation and investigation possible.