Batch-Invariant Kernel Design: Conv2d and Attention BMM

Aug 2026

The mathematical meaning of a neural-network inference call does not depend on who happens to be in the batch beside an input. The floating-point program that an accelerator runs often does. A library can select a different tile shape, reduction tree, or kernel once the batch changes. Since floating-point addition is not associative, the output for one request can move by a few ulps merely because unrelated requests arrived with it.

That is a small numerical difference, but it is not a small systems property. In autoregressive or action-generating models, a small early difference can select a different token or action and then compound. The useful contract is:

f(x₀) === f([x₀, x₁, …])[0]

Here === means bitwise equality in one fixed hardware and software environment. It is intentionally narrower than cross-device determinism, but it is exactly the guarantee a serving system needs when it changes batching.

The VLA trace exposed two missing kernels

The batch-invariant PiZero experiment traced a vision-language-action inference path while comparing a reference example alone against the same example at index zero in larger batches. Two operations were particularly instructive.

• The SigLIP patch embedder uses a 14 × 14 convolution. This is a large reduction over image channels and patch elements, and regular CUDA convolution is free to schedule that reduction from the full input shape.
• Attention multiplies rank-3 tensors after flattening batch and heads. Although the source says torch.matmul, the dispatched operator isaten::bmm, not aten::mm. An mm-only replacement therefore leaves this attention path untouched.

I 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 amount of unrelated work in the batch.

Attention: replace the operator that actually runs

A typical attention score calculation has shapes like this after the leading dimensions are flattened:

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

These are rank-3 operands, so PyTorch lowers the product to batched matrix multiplication. The new BMM kernel uses a two-dimensional launch: one dimension is the persistent tile workers and the other identifies a single BMM element. Each element gets its own base pointers and the same ordered reduction overK. Increasing B adds independent grids; it does not turn an existing sample into a different-shaped GEMM.

This distinction is the whole design. Replacing a friendly-looking Python API is not enough. The replacement must sit at the ATen operator that the shape actually selects. The demo records a profiler trace and asserts that this rank-3torch.matmul contains aten::bmm.

Conv2d: isolate each sample's reduction

The convolution path uses an explicit unfold then GEMM construction. For one image, unfold creates a matrix of patches; a flattened convolution filter matrix multiplies it, and bias is added afterward. The important implementation decision is not the algebra—it is the loop over samples. Each sample receives its own unfold and its own fixed-shape GEMM.

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

This is the generalization of the PiZero SigLIP reference implementation. The library implementation is registered for aten::convolution and supports regular non-transposed 2-D convolution, including stride, padding, dilation, and groups. It is an inference-oriented path: simple, explicit, and designed for a stable reduction schedule rather than for maximum throughput.

Correctness is two tests, not one

A replacement kernel needs to satisfy both numerical equivalence and the stronger batching property. The new CUDA demonstration covers both the 14×14 Conv2d and attention-shaped BMM cases:

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

Standard PyTorch is reported rather than unconditionally asserted to differ. The difference depends on GPU, driver, PyTorch version, and the selected kernel; an exact match on one run is not a guarantee that a later serving configuration will preserve it. The batch-invariant path is the assertion because its reduction structure is fixed by construction.

What this does—and does not—claim

These kernels do not make an entire VLA bitwise deterministic. Softmax, norms, compiler fusion, cache operations, sampling, and every other reduction still need to be traced and checked individually. Nor is per-sample Conv2d automatically the best performance choice. It trades some batching efficiency for a useful isolation guarantee.

The broader lesson is practical: batch invariance should be treated as an operator-level property. 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 its neighbors. That turns a frustrating source of serving drift into a concrete, testable systems contract.

The accompanying batch-invariant PiZero experiment contains the traced VLA reference and the original per-sample patch-projection implementation.