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

CHANGELOG.md

@@ -22,6 +22,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ### Changed
 
+- The eager forward pass no longer rebuilds the sine position embedding, the per-level padding masks, or the padded batch tensor on every call when the batch carries no padding. `nested_tensor_from_tensor_list` already decides, from Python-side shapes alone, whether any image needed padding; it now records that on `NestedTensor.no_padding` (default `False`, so every other construction site keeps the previous conservative behavior) instead of discarding it. For a batch flagged that way the padded tensor is the input itself rather than a fresh allocation plus a full copy, each feature level's mask is built directly as zeros instead of interpolating an all-False mask, and `PositionEmbeddingSine` serves the embedding from a per-`(shape, device, layout)` cache bounded to 8 entries. That cache cannot go stale: the module holds no parameters and no buffers, so for an all-False mask the embedding is a pure function of the mask's shape and device. Batches that carry real padding, including divisor round-up from `block_size`, are untouched and still read the mask. Outputs are bit-identical.
+
 - `RFDETR.predict()` now converts PIL and uint8 NumPy inputs from HWC byte storage to contiguous CHW floating-point storage with one dtype/layout allocation and in-place scaling. The default source-image path reuses its already-materialized PIL array, while non-uint8 NumPy inputs retain torchvision's conversion path. The `[0, 1]` range-scan skip for those same two input types is unaffected: the fused conversion divides `uint8` storage by 255 and carries the same guarantee `to_tensor` did.
 
 - The deformable-attention core now reuses its sampled tensor directly for one-level inputs instead of stacking and flattening a one-element list. Multi-level packing is unchanged. This primarily removes allocation work from keypoint cross-attention, whose current configuration uses one feature level and many more queries than the detection decoder.

src/rfdetr/models/backbone/backbone.py

@@ -148,28 +148,44 @@ def export(self) -> None:
             logger.info("Merging and unloading LoRA weights")
             self.encoder = self.encoder.merge_and_unload()
 
+    @staticmethod
+    def _level_mask(tensor_list: NestedTensor, feat: Tensor) -> Tensor:
+        """Downsample the batch padding mask onto *feat*'s spatial grid.
+
+        Args:
+            tensor_list: Batch whose ``mask`` marks padded input pixels.
+            feat: Feature map (B, C, H, W) whose grid the mask is resampled onto.
+
+        Returns:
+            Boolean mask of shape (B, H, W), True on padded positions.
+        """
+        if tensor_list.no_padding:
+            # Nearest-neighbour resampling of an all-False mask is all-False at every output
+            # size, so the interpolation below is a constant of ``feat``'s shape here. Same
+            # substitution ``forward_export`` already makes under its no-padding assumption.
+            batch, _, height, width = feat.shape
+            return torch.zeros((batch, height, width), dtype=torch.bool, device=feat.device)
+        m = tensor_list.mask
+        assert m is not None
+        return F.interpolate(m[None].float(), size=feat.shape[-2:]).to(torch.bool)[0]
+
     def forward(self, tensor_list: NestedTensor) -> tuple[list[NestedTensor], list[NestedTensor] | None]:
         """"""
         # (H, W, B, C)
         raw_feats = self.encoder(tensor_list.tensors)
         feats = self.projector(raw_feats)
         # x: [(B, C, H, W)]
+        no_padding = tensor_list.no_padding
         out = []
         for feat in feats:
-            m = tensor_list.mask
-            assert m is not None
-            mask = F.interpolate(m[None].float(), size=feat.shape[-2:]).to(torch.bool)[0]
-            out.append(NestedTensor(feat, mask))
+            out.append(NestedTensor(feat, self._level_mask(tensor_list, feat), no_padding))
 
         cross_attn_out = None
         if self.cross_attn_projector is not None:
             cross_attn_out = []
             cross_attn_feats = self.cross_attn_projector(raw_feats)
             for feat in cross_attn_feats:
-                m = tensor_list.mask
-                assert m is not None
-                mask = F.interpolate(m[None].float(), size=feat.shape[-2:]).to(torch.bool)[0]
-                cross_attn_out.append(NestedTensor(feat, mask))
+                cross_attn_out.append(NestedTensor(feat, self._level_mask(tensor_list, feat), no_padding))
 
         return out, cross_attn_out
 

src/rfdetr/models/position_encoding.py

@@ -23,6 +23,10 @@
 
 from rfdetr.utilities.tensors import NestedTensor
 
+# Distinct (shape, device, align_dim_orders) embeddings retained per module. Single-resolution
+# inference needs one; multi-scale training cycles through a small fixed set of scales.
+_POS_CACHE_MAX_ENTRIES = 8
+
 
 class PositionEmbeddingSine(nn.Module):
     """This is a more standard version of the position embedding, very similar to the one used by the Attention is all
@@ -45,13 +49,42 @@ def __init__(
             scale = 2 * math.pi
         self.scale = scale
         self._export = False
+        # (mask shape, device, align_dim_orders) -> position embedding. Only populated for
+        # batches flagged ``no_padding``; see ``forward``. Kept in ``__dict__`` rather than as a
+        # buffer so it stays out of ``state_dict`` and is never moved by ``.to()`` — the device is
+        # part of the key instead.
+        self._pos_cache: dict[tuple[tuple[int, ...], torch.device, bool], Tensor] = {}
 
     def export(self) -> None:
         self._export = True
         self._forward_origin = self.forward
         self.forward = self.forward_export  # type: ignore[method-assign,assignment]
 
     def forward(self, tensor_list: NestedTensor, align_dim_orders: bool = True) -> Tensor:
+        mask = tensor_list.mask
+        assert mask is not None
+        if tensor_list.no_padding:
+            # With an all-False mask the embedding below is a pure function of the mask's shape,
+            # device and align_dim_orders (this module holds no parameters or buffers, and
+            # scale/temperature/normalize are fixed at construction and never reassigned by any
+            # caller in this codebase), so a cache hit is safe as long as every consumer treats
+            # the returned tensor as read-only -- true for every current caller, which only reads
+            # it via out-of-place ops (with_pos_embed's `tensor + pos`, flatten/transpose views).
+            # Recomputing it rebuilds the same ~20-op chain (cumsum, arange, pow, sin/cos, stack,
+            # flatten, cat, permute) on every forward pass.
+            key = (tuple(mask.shape), mask.device, align_dim_orders)
+            cached = self._pos_cache.get(key)
+            if cached is None:
+                cached = self._compute(tensor_list, align_dim_orders)
+                if len(self._pos_cache) >= _POS_CACHE_MAX_ENTRIES:
+                    # Bound the retained device memory; multi-scale training cycles through a
+                    # handful of shapes, so evicting the oldest keeps the working set resident.
+                    del self._pos_cache[next(iter(self._pos_cache))]
+                self._pos_cache[key] = cached
+            return cached
+        return self._compute(tensor_list, align_dim_orders)
+
+    def _compute(self, tensor_list: NestedTensor, align_dim_orders: bool) -> Tensor:
         x = tensor_list.tensors
         mask = tensor_list.mask
         assert mask is not None

src/rfdetr/utilities/tensors.py

@@ -66,9 +66,15 @@ class NestedTensor:
     Stores both the padded tensor and a boolean mask indicating padding positions.
     """
 
-    def __init__(self, tensors: Tensor, mask: Tensor | None) -> None:
+    def __init__(self, tensors: Tensor, mask: Tensor | None, no_padding: bool = False) -> None:
         self.tensors = tensors
         self.mask = mask
+        # Structural claim, not a claim about ``mask``'s contents: True only when the batch was
+        # assembled from images that already filled the padded extent, so ``mask`` is all-False.
+        # Consumers use it to treat mask-derived values as constants of the batch shape without
+        # reading device memory (which would force a device-to-host sync). Defaults to False so
+        # every other construction site keeps the conservative behaviour.
+        self.no_padding = no_padding
 
     def to(self, device: torch.device, **kwargs: Any) -> "NestedTensor":
         """Move tensors and mask to *device*.
@@ -86,7 +92,7 @@ def to(self, device: torch.device, **kwargs: Any) -> "NestedTensor":
             cast_mask = mask.to(device, **kwargs)
         else:
             cast_mask = None
-        return NestedTensor(cast_tensor, cast_mask)
+        return NestedTensor(cast_tensor, cast_mask, self.no_padding)
 
     def pin_memory(self) -> "NestedTensor":
         """Pin tensor and mask memory for faster CPU→GPU transfer.
@@ -97,6 +103,7 @@ def pin_memory(self) -> "NestedTensor":
         return NestedTensor(
             self.tensors.pin_memory(),
             self.mask.pin_memory() if self.mask is not None else None,
+            self.no_padding,
         )
 
     def decompose(self) -> tuple[Tensor, Tensor | None]:
@@ -126,7 +133,10 @@ def nested_tensor_from_tensor_list(
 
     Returns:
         NestedTensor with all images padded to the maximum spatial dimensions (rounded up to *block_size* when
-        provided).
+        provided). When *tensor_list* is already a contiguous 4-D batch that needs no padding, the returned
+        ``.tensors`` is *tensor_list* itself (same storage, not a copy): mutating it after this call also mutates
+        the NestedTensor. A non-contiguous 4-D batch instead gets an independent, contiguous copy, like every
+        other input shape.
     """
     # TODO make this more general
     if tensor_list[0].ndim == 3:
@@ -145,14 +155,33 @@ def nested_tensor_from_tensor_list(
         b, c, h, w = batch_shape
         dtype = tensor_list[0].dtype
         device = tensor_list[0].device
+        # Every image already fills the padded extent whenever no image is smaller than the
+        # per-axis maximum and block-size rounding added nothing. That is decided here from
+        # Python-side shapes alone, so the resulting all-False mask is known without reading it.
+        no_padding = all(list(img.shape) == max_size for img in tensor_list)
+        if no_padding and isinstance(tensor_list, Tensor) and tensor_list.is_contiguous():
+            # The batch tensor already *is* the padded batch: torch.zeros(batch_shape) followed by
+            # a full-extent copy_ of every image reproduces it element for element. Skip the
+            # allocation and the copy. Returned ``.tensors`` is the caller's own tensor (same
+            # storage, same strides) rather than an independent copy -- restricted to the
+            # already-contiguous case so the result always matches the layout the allocate-and-copy
+            # path would have produced. A non-contiguous input falls through to that path instead,
+            # which normalizes both layout and ownership like every other branch.
+            mask = torch.zeros((b, h, w), dtype=torch.bool, device=device)
+            return NestedTensor(tensor_list, mask, True)
         tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
-        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
+        mask = (
+            torch.zeros((b, h, w), dtype=torch.bool, device=device)
+            if no_padding
+            else torch.ones((b, h, w), dtype=torch.bool, device=device)
+        )
         for img, pad_img, m in zip(tensor_list, tensor, mask):
             pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
-            m[: img.shape[1], : img.shape[2]] = False
+            if not no_padding:
+                m[: img.shape[1], : img.shape[2]] = False
     else:
         raise ValueError("not supported")
-    return NestedTensor(tensor, mask)
+    return NestedTensor(tensor, mask, no_padding)
 
 
 # _onnx_nested_tensor_from_tensor_list() is an implementation of

tests/models/backbone/test_backbone.py

@@ -3,15 +3,17 @@
 # Copyright (c) 2025 Roboflow. All Rights Reserved.
 # Licensed under the Apache License, Version 2.0 [see LICENSE for details]
 # ------------------------------------------------------------------------
-"""Tests for dual-projector backbone joiner routing."""
+"""Tests for dual-projector backbone joiner routing and per-level mask construction."""
 
 from __future__ import annotations
 
 import torch
+import torch.nn.functional as F  # noqa: N812
 from torch import nn
 
 from rfdetr.models.backbone import Joiner
-from rfdetr.utilities.tensors import NestedTensor
+from rfdetr.models.backbone.backbone import Backbone
+from rfdetr.utilities.tensors import NestedTensor, nested_tensor_from_tensor_list
 
 
 class _FakeBackbone(nn.Module):
@@ -138,3 +140,79 @@ def test_joiner_forward_export_contract() -> None:
     assert poss[0].shape == exported_features[0][:, :1, :, :].shape
     assert isinstance(cross_attention, list)
     assert all(isinstance(feature, torch.Tensor) for feature in cross_attention)
+
+
+class TestBackboneLevelMask:
+    """``Backbone._level_mask`` substitutes zeros only where that is exactly the interpolation.
+
+    Nearest-neighbour resampling of an all-False mask is all-False at every output size, so the substitution must agree
+    with ``F.interpolate`` for unpadded batches and must not be taken when the batch carries real padding.
+    """
+
+    @staticmethod
+    def _interpolated(mask: torch.Tensor, feat: torch.Tensor) -> torch.Tensor:
+        """Reference result: the interpolation the unflagged path performs.
+
+        Examples:
+            >>> mask = torch.zeros(1, 8, 8, dtype=torch.bool)
+            >>> TestBackboneLevelMask._interpolated(mask, torch.rand(1, 4, 2, 2)).shape
+            torch.Size([1, 2, 2])
+        """
+        return F.interpolate(mask[None].float(), size=feat.shape[-2:]).to(torch.bool)[0]
+
+    def test_unpadded_matches_interpolation(self) -> None:
+        """The zeros shortcut equals what interpolating the all-False mask produces."""
+        mask = torch.zeros(2, 64, 64, dtype=torch.bool)
+        feat = torch.rand(2, 8, 16, 16)
+        flagged = NestedTensor(torch.rand(2, 3, 64, 64), mask, True)
+        assert torch.equal(Backbone._level_mask(flagged, feat), self._interpolated(mask, feat))
+
+    def test_unflagged_all_false_mask_takes_the_interpolation(self) -> None:
+        """Without the flag the mask is still read, and the result is unchanged."""
+        mask = torch.zeros(2, 64, 64, dtype=torch.bool)
+        feat = torch.rand(2, 8, 16, 16)
+        unflagged = NestedTensor(torch.rand(2, 3, 64, 64), mask, False)
+        assert torch.equal(Backbone._level_mask(unflagged, feat), self._interpolated(mask, feat))
+
+    def test_padded_mask_is_preserved(self) -> None:
+        """Real padding must survive downsampling rather than be zeroed away."""
+        mask = torch.zeros(1, 64, 64, dtype=torch.bool)
+        mask[:, 32:, :] = True
+        feat = torch.rand(1, 8, 16, 16)
+        unflagged = NestedTensor(torch.rand(1, 3, 64, 64), mask, False)
+        level = Backbone._level_mask(unflagged, feat)
+        assert torch.equal(level, self._interpolated(mask, feat))
+        assert level[:, 8:, :].all().item() is True
+        assert level[:, :8, :].any().item() is False
+
+    def test_output_shape_and_dtype_follow_the_feature_map(self) -> None:
+        """The shortcut must produce the same shape/dtype/device contract as the interpolation."""
+        mask = torch.zeros(3, 40, 40, dtype=torch.bool)
+        feat = torch.rand(3, 8, 10, 20)
+        flagged = NestedTensor(torch.rand(3, 3, 40, 40), mask, True)
+        level = Backbone._level_mask(flagged, feat)
+        assert level.shape == (3, 10, 20)
+        assert level.dtype == torch.bool
+        assert level.device == feat.device
+
+    def test_matches_unflagged_after_default_config_batch_uniform_resize(self) -> None:
+        """The shortcut still agrees with the interpolation after the real default-training mutation.
+
+        The default training config (``square_resize_div_64=True``, ``multi_scale=True``,
+        ``do_random_resize_via_padding=False``) resizes every sample to one fixed square scale before collate, so
+        the batch is flagged. ``RFDETRLightningModule.on_train_batch_start`` then resizes the whole batch uniformly
+        to a randomly chosen scale in place, without touching ``no_padding`` (see
+        ``tests/utilities/test_tensors.py::TestNestedTensorNoPadding::test_flag_survives_inplace_batch_uniform_resize``
+        for the flag/mask invariant this relies on). The two consumers of the mask must still agree afterwards.
+        """
+        images = [torch.rand(3, 512, 512) for _ in range(2)]
+        nested = nested_tensor_from_tensor_list(images, block_size=64)
+        with torch.no_grad():
+            nested.tensors = F.interpolate(nested.tensors, size=(640, 640), mode="bilinear", align_corners=False)
+            nested.mask = (
+                F.interpolate(nested.mask.unsqueeze(1).float(), size=(640, 640), mode="nearest").squeeze(1).bool()
+            )
+
+        unflagged_twin = NestedTensor(nested.tensors, nested.mask, False)
+        feat = torch.rand(2, 8, 20, 20)
+        assert torch.equal(Backbone._level_mask(nested, feat), Backbone._level_mask(unflagged_twin, feat))