perf(postprocess): pick selected rows without materialising repeated index tensors - #1268

What does this PR do?

The three query selections in PostProcess build their gather index with repeat(), so the index tensor gets materialised at the full size of the data it selects. For masks this is an int64 tensor of shape [K, Hm, Wm] — at num_select=300 that is 21 MiB per image for rfdetr-seg-small (96×96 mask head) and 84 MiB for rfdetr-seg-2xlarge (192×192), allocated and filled on every image just to pick 300 rows.

index_select picks the same whole rows without building any index tensor, and expand on the box path replaces the materialised [B, K, 4] index with a stride-0 view. The output does not change — bit for bit (see below).

The ONNX benchmark keeps its own copy of the box selection (post_process in src/rfdetr/export/benchmark.py) with the same repeat(), so it gets the same one-line change.

Measured on CPU (torch 2.13, single thread, pinned to one core, median of interleaved A/B runs of the full PostProcess.forward):

A few notes on where this matters:

• The big win is the mask path when masks stay at head resolution (upsample_masks_to_image_size=False, the eval_masks_head_resolution cost lever), where the gather is most of the postprocess. That path runs with K=300 on every validation image.
• On the normal predict path the upsample to image size dominates, so this is time-neutral there — what it removes is the 21–84 MiB index allocation per image before the resize.
• With a score_threshold (perf(inference): skip upsampling masks the caller's threshold discards #1265) K is small and everything is already cheap, nothing changes there.
• The two-stage query selection in src/rfdetr/models/transformer.py (L381, L389) builds its gather indices with the same repeat() pattern. I left it alone on purpose: it sits in the model forward under autograd, and its indices top out at [B, K, d_model] — a fraction of the mask case. Happy to cover it here or in a follow-up if you want it.

Related Issue(s): none, follow-up of #1265 on the same function.

Type of Change

• Performance improvement (no behaviour change)

Testing

• Bit-identical outputs: old vs new forward() compared over the three heads (masks / keypoints / boxes) on CPU and CUDA — 198 output tensors compared byte for byte, including duplicated top-k query indices, Q < num_select, a threshold that keeps zero rows, and both upsample_masks_to_image_size branches. All equal.
• Three new tests pin the contract the change relies on: top-k can select the same query under two classes, and each occurrence must reproduce the source row verbatim (test_duplicate_query_selection_repeats_the_same_mask_rows, test_gather_keypoints_for_queries_repeats_duplicated_indices, test_post_process_repeats_boxes_for_duplicated_topk_queries for the benchmark copy). All pass before and after the change — they test the equivalence, not the implementation.
• The benchmark copy was verified the same way: old vs new post_process bit-compared (int32 view) over 50 randomised shapes, including k clamped below and above Q*C. All equal.
• tests/models plus tests/inference/test_trt_inference.py pass with the patch (643 passed, 5 skipped) and ruff check/ruff format are clean on the touched files .. let me know if you want numbers for any other case.