perf: stop paying kornia's per-module container cost on every batch

[Hendrik-code]

Follows #126, which has landed; this branch is off main and the diff is only
the container work. The tables below carry a #126 column so the two steps can
be told apart — it is the commit now at the tip of main.

#126 made the augmentations faster and left the container they run in alone. Set
every p to 0 so that nothing fires, and a 26-module pipeline on a
[2, 1, 128, 128, 128] batch still took 13.9 ms on origin/main and
13.0 ms after #126. That is AugmentationSequential itself, paid once per
training batch whether or not a single augmentation does anything. It is now
3.9 ms.

What it was spending it on

201 MB of segmentation, copied per call. kornia deep-copies the parameter
dict on the way into every leaf — _process_kwargs_to_params_and_flags →
override_parameters → deepcopy_dict, which clones every value that is a
Tensor, once per data key per module. SmaugLab puts the segmentation in that
dict, so the label map was duplicated 78 times a batch:

origin/main:   279 tensor clones inside deepcopy_dict,  201.3 MB
                  seg  78 clones  201.33 MB
this branch:    22 tensor clones inside deepcopy_dict,    0.0 MB

SegmentationRef carries it behind an object, which deepcopy_dict passes
through by reference, and resolves kornia's row selection on first read instead
of eagerly — most modules never look at it. Consumers read it with
segmentation_from(params), which still accepts a bare tensor from a direct
apply_transform call.

The leaf forward chain. _apply_leaf inlines forward → apply_func →
transform_inputs for this package's 3-D leaves, skipping two parameter
deep-copies, three transform_tensor round trips and four calls to kornia's
is_autocast_enabled — which re-parses torch.__version__ with
packaging.version.Version on every call: 7.7 µs, against 0.08 µs for
the builtin it wraps, roughly a hundred times a batch.

The mask pass for transforms that cannot touch a mask. Most of a pipeline is
image-only, and for those transform_masks deep-copies, reshapes, validates and
resolves the row selection in order to return the tensor it was handed.

Constants rebuilt per module. The 4×4 identity (torch.eye plus a repeat,
26× a call) and forward_input_shape (torch.tensor of a five-element shape,
8.5 µs, 27× a call) are cached on the patch shape and expanded rather than
allocated. _leaf_parameters also drops the one-element allocation, sum and
broadcast multiply that __batch_prob_generator__ performs to decide something
p_batch == 1 already fixed.

Net effect on the Python side: cProfile over 20 zero-probability calls goes from
0.276 s to 0.098 s.

Measurements

Idle NVIDIA L40S, torch 2.6.0+cu124, kornia 0.8.2, [2, 1, 128, 128, 128] with a
3-label segmentation. Median of 4 interleaved rounds of 50 calls — the three
checkouts take turns within each round, so GPU clock ramp hits all of them
equally, and each process warms up 20 calls before timing.

config	mode	main	#126	this	vs main	vs #126
gpu.json	sequential	24.66	20.45	10.61	2.32x	1.93x
gpu.json	random_order	17.91	10.98	8.83	2.03x	1.24x
gpu.json	random_order_ta	19.38	16.50	11.57	1.68x	1.43x
one-sequence-to-segment.json	sequential	21.12	19.16	9.89	2.13x	1.94x
one-sequence-to-segment.json	random_order	13.89	9.01	8.07	1.72x	1.12x
one-sequence-to-segment.json	random_order_ta	17.32	11.91	7.88	2.20x	1.51x
gpu.json [all p=0]	sequential	13.90	13.01	3.88	3.58x	3.35x
one-sequence… [all p=0]	sequential	12.69	11.43	3.54	3.58x	3.23x

The [all p=0] rows are the container on its own.

A note on comparing with #126's table: that one timed each checkout in a single
long block, which lets the GPU settle at a high clock; this one restarts a
process every ~20 s so the clock spends more time ramping. Absolute numbers are
accordingly higher here (main/sequential reads 24.66 against 18.16 there). The
ratios are what transfer between the two protocols, and within this table the
conditions are matched.

Correctness

Bit-identical output. Both shipped configs, in all three pipeline modes,
image and mask, over twelve seeds on CPU: max|diff| = 0 against #126 in every
one of the twelve series. This is a pure-overhead change and that is what a pure
overhead change should produce.

On CUDA the comparison bottoms out at ~1e-5, because the pipeline is already
non-deterministic between two runs of the same checkout — origin/main
included (1.2e-5 … 3.8e-5 between two runs of it). Float atomics; not introduced
here, and worth knowing independently.

The short-circuits are checked against kornia, not assumed. They duplicate
kornia control flow, and kornia is a moving dependency — tests.yml runs a
compat matrix over several versions. So unit_tests/test_container_fast_path.py
(22 tests, 281 subtests) runs every registered transform down both routes, for
all-apply, no-apply and the partial-batch case, and asserts that

• the output tensors match,
• the transform matrix each route stores on the module matches,
• the parameter dictionaries _leaf_parameters and kornia's
  forward_parameters draw match tensor-for-tensor, across a grid of
  p / p_batch / same_on_batch and three seeds,
• the mask short-circuit returns exactly what transform_masks returns,

plus the cases where the fast path must decline: a 4-D volume, non-empty
extra_args, a foreign module, an open autocast region, and a subclass that
overrides a mask hook.

The two private kornia helpers the parameter draw needs (_get_new_batch_shape,
ImageSequentialBase) are imported behind a try/except; a kornia that moved
them gets its own loop back, and a test exercises that fallback path.

455 passed, 3 skipped, 1596 subtests; ruff check, ruff format --check and
mypy clean.

What is left

forward_parameters is still the largest single block of the remaining ~3.9 ms:
the per-module probability draw and the geometric generators, which genuinely
have to sample. Cutting further would mean skipping parameter generation for
modules that will not fire, which changes the RNG stream under same_on_batch
and is not worth it for the ~0.4 ms on offer.

🤖 Generated with Claude Code