perf(inference): skip padding-mask work the batch already proved unnecessary - #1416

What

nested_tensor_from_tensor_list already decides, from Python-side shapes alone, whether any image in the batch needed padding. It throws that away. Downstream the all-False mask it produced is then read on every forward pass to rebuild three values that are constants of the batch shape.

This records the decision on NestedTensor.no_padding and uses it:

• nested_tensor_from_tensor_list — when the input is already a 4-D batch and nothing needed padding, the padded tensor it would build is the input element for element, so the allocation and the full copy are skipped. When the input is a list, the mask is built as zeros instead of ones plus one write per image.
• Backbone._level_mask — nearest-neighbour resampling of an all-False mask is all-False at every output size, so the per-level F.interpolate becomes torch.zeros. Backbone.forward_export already makes exactly this substitution under its own no-padding assumption.
• PositionEmbeddingSine.forward — for an all-False mask the embedding is a pure function of the mask's shape and device, so it is served from a cache bounded to 8 entries instead of rebuilding cumsum, arange, pow, sin/cos, two stack+flatten, cat and permute every call.

A batch that carries real padding — a ragged list, or block_size divisor round-up — is untouched and still reads the mask. The default is no_padding=False, so the flag is off unless nested_tensor_from_tensor_list positively established there was nothing to pad. Two branches deliberately stay unflagged and so keep the previous behaviour exactly: the ONNX-tracing branch, which returns from _onnx_nested_tensor_from_tensor_list before any of this, and RFDETRDataModule.transfer_batch_to_device, which rebuilds the batch after a Kornia geometric transform that can introduce real padding.

Why a cache hit is safe

PositionEmbeddingSine holds no parameters and no buffers (list(module.parameters()) == [], list(module.buffers()) == [], state_dict() == {} — asserted in test_module_holds_no_parameters_or_buffers). With the mask contents fixed by the flag, the remaining inputs are shape, device and layout, all of which are in the key, plus scale/temperature/normalize, which are set once in __init__ and never reassigned by any call path in this codebase. A cache hit returns the same tensor object on every call, so it is also safe only because every current consumer reads it out-of-place (with_pos_embed's tensor + pos, and the flatten(2).transpose(1, 2) view in transformer.py) rather than mutating it in place. The cache lives in __dict__, not as a buffer, so it stays out of state_dict and is never silently moved by .to() — the device is part of the key instead.

Results

Ephemeral GCP L4 (g2-standard-8, torch 2.9.1+cu129), base 6674d858. Each cell is the median of 8 alternating baseline/patched fresh processes, each reporting the median of 60 timed calls after 30 warm-ups; the bracket is the min/max across those 8 process medians.

LWDETR.forward:

Public RFDETR.predict() on one real COCO image (val2017/000000039769.jpg, 640x480), 6 alternating processes per model, 60 timed calls each:

The work removed is CPU launch overhead, which is roughly constant per forward call regardless of batch size, while GPU compute time grows with it. Base already crosses into GPU-bound territory by batch 4; nano, the smallest and fastest model, needs batch 8 — those two rows are reported to show that crossover, not a regression. Small and medium were not separately measured at batch >=4; being between nano and base in compute cost, their crossover point is expected to fall somewhere in between, not necessarily at batch 4.

Correctness

• Bit-identical outputs for nano/small/medium/base at batch 1, base at batch 4 and nano at batch 8: pred_boxes and pred_logits compared with torch.equal across separate baseline and patched processes, max absolute difference 0.0 in every configuration.
• predict() returns the same detection count and the same coordinate and confidence checksums on every one of the 36 measured runs.
• The benchmark re-runs 20 forwards and asserts each equals the first, so a consumer mutating the cached embedding in place would fail rather than pass silently.
• 16 of the new tests are red on a pristine checkout of the same base commit and green with the patch (plus test_position_encoding.py, which fails to import _POS_CACHE_MAX_ENTRIES before the patch).

Not covered

• I did not re-run COCO val2017 mAP. Outputs are bit-identical, which is byte parity, not a separately measured mAP number — happy to run it if that's wanted before merging.
• I only measured on an L4, torch 2.9.1+cu129. The saving is CPU-side launch overhead, so its size depends on the host/GPU ratio.
• The 8-entry position-embedding cache retains device memory that scales with batch size, and that number doesn't come from a measured training run. Simulating RFDETRSmall's default multi-scale schedule (resolution=512, expanded_scales=True) filling all 8 cache slots retains 30.94 MiB at batch 4 and 123.75 MiB at batch 16 (repro/pos_cache_memory.py). That is bounded — it cannot grow past 8 entries — but it is not free, and a training setup already close to its VRAM ceiling should account for it.
• I did not time a training loop, and the flag is not simply "off" during training by default. The Kornia GPU-augmentation path is confirmed unflagged: transfer_batch_to_device rebuilds the batch as NestedTensor(img_aug, padding_mask_aug) after a geometric transform that can introduce real padding. But that path is opt-in (augmentation_backend="kornia"/"auto" with CUDA); the default augmentation_backend="cpu" resizes every sample to one fixed square scale before collate (square_resize_div_64=True, multi_scale=True, do_random_resize_via_padding=False all default), so those batches collate with the flag set. The per-batch multi-scale step that follows (on_train_batch_start) then resizes the whole batch uniformly in place without touching the flag; a new test confirms it still describes the mask correctly afterwards and that the backbone mask and position-embedding cache match the unflagged path bit for bit on that mutated batch. do_random_resize_via_padding=True and COCO validation's variable-sized images stay correctly unflagged for the same reason as Kornia. I don't know if any of this moves measured step time .. the backward pass, optimizer step, and data loading may dwarf the forward-only overhead measured above.