Batch-Invariant VLAs

Aug 2026

Not many people care about VLA inference servers, but I do. I care because in the very near future they'll be powering fleets of robots both in production and during training (RL).

Non-determinism in VLA inference is expected since it has already been shown to exist in LLM inference. However, the non-determinism in VLA inference is more insidious because it can be propagated through the action space and affect the robot's behavior in the real world.

This post describes a debugging journey to make π0 batch-invariant, which is crucial for reliable robot behavior.

π0 is a vision-language-action policy: combines images and robot state with a language-conditioned model, then predicts a chunk of continuous actions through flow matching.

Batch-invariant: batching several requests together must not change the answer for any one of them.

TL;DR: a per-sample SigLIP patch projection and an explicit stack of two-dimensional torch.mm calls were needed in the π0 path. The latter was a way to force the dispatcher to reach the invariant aten::mm implementation (provided by Thinking Machines) instead of silently taking a batched path.

Profiling batch-invariance

Before changing any code, I added an opt-in tracer and began recording named tensors along the model path:

with trace_context(trace):
    output = model(**inputs)

GLOBAL_TRACE.record("final_action", action)

The tracer recursively detaches and clones tensors, including tensors nested in dictionaries and lists. It moves those snapshots to CPU, then compares the batch-one reference with element zero of the batched run. At each point it reports maximum absolute difference, mean absolute difference, RMSE, relative L2 error, and output magnitude.

"input_embeds.selected_image_feature.hidden_states": {
  "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": 4.864709377288818
},
"input_embeds.selected_image_feature.last_hidden_state": {
  "max_abs_diff": 6.496906280517578e-06,
  "mean_abs_diff": 1.1066537126680487e-06,
  "rmse": 1.3891448134017992e-06,
  "relative_l2_mean": 8.866771281645924e-07,
  "relative_l2_max": 8.866771281645924e-07,
  "output_abs_max": 7.3877058029174805
},
"input_embeds.selected_image_feature.last_hidden_state.post_layernorm": {
  "max_abs_diff": 4.26173210144043e-06,
  "mean_abs_diff": 7.08217100964248e-07,
  "rmse": 8.896561212168308e-07,
  "relative_l2_mean": 8.896386702872405e-07,
  "relative_l2_max": 8.896386702872405e-07,
  "output_abs_max": 4.708508491516113
},

The names preserve the route through the model. For example, the image path records names such asinput_embeds.selected_image_feature.patch_embeds and and joint-model states understep_0.action_joint_model.layer_<n>.pre_attn and...final. This allows us to find the first diverging (non-zero diff) tensor.

Input
  ↓
SigLIP patch embedding        exact
  ↓
vision transformer            exact
  ↓
joint model / KV cache        exact
  ↓
action expert                 exact
  ↓
flow integration              exact
  ↓
final action                  exact

After a fix, the same experiment can be rerun. If the earlier tensors are now exact but a later tensor is not, the tracer has done its job: it has moved the investigation to the next actionable boundary.

Thinky's Ops were not enough

Before changing π0, I enabled the batch-invariant operator mode from Thinking Machines. Their work already provides CUDA implementations whose arithmetic is designed not to change when unrelated batch members are added. In this checkout the mode uses torch.library and registers:

• aten::mm
• aten::addmm
• aten::_log_softmax
• aten::mean.dim

The important entries for this story are aten::mm andaten::addmm. Their implementations call a persistent Triton matmul kernel, accumulate output tiles in float32, and use a fixed tile configuration selected by input dtype.

I reran the same batch experiment. Some differences disappeared, but the action was still not invariant. The tracer painted a picture like the following:

Input
  ↓
SigLIP patch embedding        diverged
  ↓
vision transformer            diverged
  ↓
joint model / KV cache        diverged
  ↓
action expert                 diverged
  ↓
flow integration              diverged
  ↓
final action                  diverged

Yay! I was worried that Thinky's batch-invariant operators were going to be enough and that the project would be easy.

Weird nn.Conv2d Behavior

The tracer pointed to the very beginning of π0's visual path: SigLIP's patch embedding. This is the operation that turns an image into tokens. In the configured model, a 224×224 RGB image is split into non-overlapping 14×14 patches, producing(224 / 14)² = 256 image tokens. The original path usednn.Conv2d with 1,152 output channels, a 14×14 kernel, a 14×14 stride, and no padding.

Mathematically, each output patch is independent of every other sample:

y[b, p, o] = Σ x[b, c, hₚ + kₕ, wₚ + k𝓌] · W[o, c, kₕ, k𝓌]

However, the common implementations of any matmul (which includes convolution) uses a reduction over the input channels and kernel positions. The order of that reduction can change when the global shape changes, which can lead to non-deterministic behavior across batches.

After testing the nn.Conv2d kernel on its own, with the same configuration, we found out that it stops being batch-invariant at batch-size 32.

Conv2d(
  in_channels=3,
  out_channels=1152,
  kernel_size=(14, 14),
  stride=(14, 14), padding=(0, 0),
  padding_mode='zeros', device='cuda')

batch_size =  1: max abs diff = 0.0
batch_size =  2: max abs diff = 0.0
batch_size =  4: max abs diff = 0.0
batch_size =  8: max abs diff = 0.0
batch_size = 16: max abs diff = 0.0
batch_size = 32: max abs diff = 0.000902771949...

On paper, this should be an easy case. Every output patch depends only on one image, one patch, and one set of weights. But that mathematical independence does not force one fixed GPU reduction order. The convolution implementation can make choices based on the complete input shape.

Replacing nn.Conv2d

The active SigLIP model now uses _UnfoldConv2d: for each sample it callsF.unfold, flattens the learned kernel, computes a matrix product, adds the bias, and reshapes the result back to(1, channels, patch_height, patch_width) before concatenating samples. It is intentionally straightforward. I wanted the reduction scope to be visible before worrying about performance.

for sample in x:
    patches = F.unfold(sample.unsqueeze(0), ...)
    output = flattened_weight @ patches + bias[:, None]
    outputs.append(output.reshape(1, -1, height, width))
return torch.cat(outputs, dim=0)

What about a custom Triton implementation?

I also explored writing a specialized Triton implementation. It maps each program to an output element, decodes width, height, channel, and sample indices, then performs a fixed reduction over input channels and kernel positions. It masks out-of-range reduction elements and padded spatial loads, accumulates in a Triton accumulator, adds bias, and stores one output element.

The design process is out of scope for this investigation and will be explored in a future post.

torch.matmul(a,b) != torch.mm(a,b)

The new _UnfoldConv2d implementation successfully restored batch-invariance to the SigLIP patch projection. However, the tracer then pointed to a new divergence in the joint model, in the attention path.

Input
  ↓
SigLIP patch embedding        exact
  ↓
vision transformer            exact
  ↓
joint model / KV cache        diverged
  ↓
action expert                 diverged
  ↓
flow integration              diverged
  ↓
final action                  diverged

In this case there was no smell. All operators used the batch-invariant package (aten::mm, aten::addmm, aten::_log_softmax, and aten::meandim). So the error must be in the dispatch path.

Adding more internal logging steps for the tracer to check it seems like the divergence begins at the final matrix multiplication of the joing model.

attn_output = torch.matmul(attn_weights, value_states)

GLOBAL_TRACE.record(f"step_0.action_joint_model.layer_{layer_idx}.pre_attn.attn_output_pre", attn_output)

Wait: shouldn't Thinky's batch-invariant aten::mm have been used here? The answer is no.
torch.matmul is shape-dependent: it does not always dispatch to the same operator as torch.mm.

With two-dimensional inputs, torch.matmul reachesaten::mm. With higher-rank inputs, it treats the leading dimensions as batches and reaches aten::bmm. The Python call looks similar, but the profiler shows a different dispatch path:

# attn_weights = torch.randn(4, 281, device='cuda')
# value_states = torch.randn(281, 256, device='cuda')
# attn_output = torch.matmul(attn_weights, value_states)

ProfilerResult:
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                           aten::matmul         0.05%      21.280us        99.99%      46.617ms      46.617ms       0.000us         0.00%      31.353us      31.353us             1  
                                               aten::mm        76.89%      35.848ms        99.94%      46.596ms      46.596ms       4.479us       100.00%      31.353us      31.353us             1  

---------------

# attn_weights = torch.randn(32, 8, 4, 281, device='cuda')
# value_states = torch.randn(32, 8, 281, 256, device='cuda')
# attn_output = torch.matmul(attn_weights, value_states)

ProfilerResult:
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                                   Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg     Self CUDA   Self CUDA %    CUDA total  CUDA time avg    # of Calls  
-------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  
                                           aten::matmul         0.14%      64.853us        99.99%      47.305ms      47.305ms       0.000us         0.00%     211.712us     211.712us             1  
                                              aten::bmm        75.93%      35.922ms        99.69%      47.166ms      47.166ms      26.464us       100.00%     211.712us     211.712us             1  
            

Replacing aten::bmm

Simple Looped Solution

A simple looped solution would be to treat the first two dimensions as independent batch-like indices and performs a genuine two-dimensional multiplication for every pair:

def loop_bmm(a, b):
    return torch.stack([
        torch.stack([
            torch.mm(a[i][j], b[i][j])
            for j in range(a.shape[1])
        ])
        for i in range(a.shape[0])
    ])

Using vmap

Using vmap(torch.mm) sounds promising. Until you realise that vmaprequires special registration using torch.library.register_vmap.

That basically means that vmap is not a magic solution that will automatically work for all functions. It needs to be explicitly implmented for our new aten::mm implementation. However, if we do torch.vmap(torch.mm)(a, b) anyway, it would simply dispatch the existing aten::bmm implementation, which is not what we want.

A custom batch-invariant aten:bmm implementation for vmap is explored in a future blog.

Using a custom Triton Implementation

A custom Triton implementation for aten::bmm is also explored in a future blog.

Running the tracer after all our changes

Conclusion

The investigation started with Thinking Machines' batch-invariant operators, but enabling them was only the first experiment. The tracer then gave me a sequence of concrete failures instead of one mysterious final-action mismatch: first the SigLIP patch projection, then the higher-rank attention product.

The patch-projection replacement made the first trace segment exact. Then the next mismatch appeared in higher-rank attention multiplication. The Python spelling torch.matmul did not imply the replaced aten::mm path, andvmap(torch.mm) was not a safe assumption either. The explicit stack of torch.mm calls was the simplest way to force the invariant implementation and test the dispatch diagnosis.

That left me with a correctness-first implementation and a clear performance debt. The remaining optimization is to replace the launch-heavy stack loop with a batch-invariant BMM kernel that preserves the same arithmetic contract. For me, that is the useful end state of this kind of systems work: find the first divergence, prove the dispatch path, restore the numerical contract, and only then make the path fast.

Numerical guarantees in PyTorch live below the Python API. If the guarantee depends on a specific arithmetic path, profiling and dispatch inspection are not optional.

Attribution

The batch-invariant operator approach is based on the work published by Thinking Machines. The model is derived from the open-pi-zero implementation and the π0 architecture described by Physical Intelligence. The implementation uses PyTorch and Triton. The π0-specific work here is the first-divergence tracing, the SigLIP patch-projection intervention, the dispatch investigation, the explicit BMM workaround, and the end-to-end experiment harness.