perf(rfdetr-seg): Triton fusion + CUDA graphs + scratch caching (109 → 151 FPS on vehicles_312px)

[aseembits93]

Summary

RF-DETR nano seg TensorRT e2e latency optimizations, benchmarked on vehicles_312px.mp4 (538 frames, 312×176) with the minimal InferencePipeline benchmark script.

Metric	Baseline	Current	Δ
Avg FPS (Tesla T4, FP16 TRT, vehicles_312px.mp4)	109	~151	+38%
Parity vs v16 best	—	0 diff across 538 frames	bit-exact

All new paths are opt-in behind env vars; defaults preserve current behavior.

Benchmark:

python development/stream_interface/rfdetr_nano_seg_trt_workflow.py \
    --video_reference /home/ubuntu/inference/vehicles_312px.mp4

Flags to enable:

• RFDETR_USE_TRITON_PREPROC=true
• RFDETR_TRITON_FULLPOSTPROC=true
• ENABLE_AUTO_CUDA_GRAPHS_FOR_TRT_BACKEND=true

---

Commit-by-commit

#	Commit	Summary
1	419c5a7 initial point	Branch starting point.
2	9ec6a50 fused Triton pre/postproc kernels, GPU mask decode, stream sync reduction	First landing: Triton preprocess kernel, Triton conf-filter postprocess, CUDA-graph enablement, drop pre/post stream syncs.
3	807ea1a replace benchmark script with minimal InferencePipeline-based version	Switch benchmark harness to a 116-line InferencePipeline driver so measurements match production code paths.
4	5ff27d1 full Triton postproc fusion + deferred counter sync (W2)	Collapses the post-TRT chain into two Triton kernels (_rfdetr_fullpost_filter_kernel + _rfdetr_fullpost_mask_kernel_compact). Removes the CPU-blocking counter.item() in the Triton wrapper; counter is now pinned-DtoH'd from the adapter under a torch.cuda.Event guard. Includes atomic-counter compaction, _trt_reuse_as_input_buffer marker (eliminates per-frame DtoD from preproc output → graph input buffer), pinned async DtoH, cross-stream event plumbing (no more stream.synchronize() in the graph-replay branch).
5	c1406a8 cache per-call scratch + class_mapping int32 view (W7)	Eliminates two stray per-frame kernel launches that nsys flagged: direct_copy_kernel_cuda (from class_mapping.to(dtype=torch.int32) every frame — now cached by id()) and the three torch.empty allocator calls per frame (now a persistent scratch cache keyed on (num_queries, device)). counter.zero_() still fires FillFunctor per frame (safe to inline into the filter kernel is not possible — concurrent blocks would race with the zero).

---

What's in

New Triton kernels

• inference_models/models/rfdetr/triton_preprocess.py — fused stretch-to resize + BGR→RGB + /255 + ImageNet normalize. Replaces ~8 torch CUDA kernels with 1.
• inference_models/models/rfdetr/triton_postprocess.py — fused sigmoid + argmax-over-classes + class-remap + confidence-threshold filter. Used when RFDETR_TRITON_POSTPROC=true alone.
• inference_models/models/rfdetr/triton_fullpostproc.py (new in W2) — full post-process fusion:
  • _rfdetr_fullpost_filter_kernel — per-query: sigmoid argmax + class remap + conf threshold + cxcywh→xyxy + letterbox-denormalize + clip + banker's rounding; tl.atomic_add(counter, 1) reserves a compact output slot.
  • _rfdetr_fullpost_mask_kernel_compact — GPU-side bilinear upsample 78×78 → orig_h × orig_w + threshold > 0 + uint8 emit. Reads counter on GPU for early-exit so no CPU sync between the two launches.
  • Returns unsliced buffers + counter + done_event; adapter does the host-side slice after a single pinned counter DtoH.
• inference/models/rfdetr/triton_preprocess.py — same preprocess kernel wired into the legacy inference/models/rfdetr path for parity.

Wiring / dispatch

• inference_models/models/rfdetr/rfdetr_instance_segmentation_trt.py
  • Triton preprocess fast-path in pre_process() with strict guard (STRETCH_TO, numpy HWC BGR uint8, no static crop).
  • Pre-allocated input buffer with _trt_reuse_as_input_buffer marker — graph capture bakes this tensor's address as the engine's input buffer, eliminating the per-frame 1.1 MB DtoD copy.
  • Pinned host staging buffer (_fast_src_host_pinned) so copy_(non_blocking=True) actually runs async.
  • Cross-stream event recording (_fast_preproc_event) instead of pre_process_stream.synchronize().
• inference_models/models/rfdetr/common.py — dispatches RFDETR_TRITON_FULLPOSTPROC first (batch=1, no static crop, no nonsquare-intermediate resize, class remapping active), falls back to Triton conf-filter, then to torch.
• inference_models/models/common/trt.py:
  • TRTCudaGraphState.consumer_done_event lets the next graph replay wait on the consumer's last use of the output buffers (avoids DtoD clones of output buffers).
  • _capture_cuda_graph(use_pre_processed_images_as_input_buffer=...) supports capturing with an externally-owned input tensor.
  • Removed stream.synchronize() in infer_from_trt_engine's graph-replay branch; replaced with a produce_event recorded on the graph's own stream.
• inference/core/models/inference_models_adapters.py — GPU fast-path for the _combined_gpu / _counter_gpu / _postproc_done_event side channels: waits the done_event, pinned-DtoH's the 4-byte counter, syncs once to read n_survivors, then slices combined[:n] + mask[:n] and pinned-DtoH's both async with a second sync. Bitcasts column 4 of the combined int32 buffer to fp32 via numpy.view(np.float32) for confidence.

Scratch caching (W7)

• _SCRATCH_CACHE — combined/survivor_idx/mask_any/counter tensors reused across frames, keyed by (num_queries, device).
• _CLASS_MAPPING_INT32_CACHE — int32 view of the class_mapping tensor cached by id(source_tensor); previously re-converted every frame because upstream stores it as int64.

Benchmark harness

• development/stream_interface/rfdetr_nano_seg_trt_workflow.py — replaced with a minimal 116-line InferencePipeline-based driver. Uses the roboflow_core/roboflow_instance_segmentation_model@v3 workflow block; single video source, single sink, no annotators/buffering/rate-limiting. --backend flag pre-import sets DISABLED_INFERENCE_MODELS_BACKENDS for unambiguous backend attribution.

Env vars (all opt-in)

Var	Default	Effect
RFDETR_USE_TRITON_PREPROC	off	Enable fused Triton preprocess kernel
RFDETR_TRITON_POSTPROC	off	Enable Triton conf-filter (partial postproc fusion)
RFDETR_TRITON_FULLPOSTPROC	off	Enable full postproc fusion (W2)
RFDETR_GPU_POSTPROCESS	on	GPU mask-emptiness skip
RFDETR_DISABLE_GPU_PREPROC	off	Opt-out for torch GPU preprocess branch
ENABLE_AUTO_CUDA_GRAPHS_FOR_TRT_BACKEND	off	Enables TRTCudaGraphCache

---

Test plan

• pytest tests/inference/unit_tests/models/test_rfdetr.py -x -q — 11/11 pass
• Benchmark on vehicles_312px.mp4 (538 frames, Tesla T4, FP16): ~151 FPS (baseline ~109)
• Parity check on vehicles_312px — 0 diff across all 538 frames vs v16 best (bit-exact xyxy, confidence, class_id, per-detection mask MD5)
• nsys: kernel count 170,530 → 2,931 (-98.3%) with CUDA graphs + full postproc fusion enabled; DtoD per-run 439 → 8 with _trt_reuse_as_input_buffer
• Tracer confirms zero class_mapping.to() / direct_copy_kernel_cuda launches per frame after W7
• Higher-resolution video (mask kernel compute scales with orig_h × orig_w)
• Jetson Orin NX (pinned-memory + async-transfer wins should scale better)

What closed out

Supersedes #23 (same work, earlier snapshot).

Known remaining bubbles

From nsys on vehicles_best:

• cudaStreamSynchronize is 80.9% of CUDA API time — CPU waits for GPU, GPU is not saturated. The serial per-frame dispatch is the next bottleneck.
• _rfdetr_fullpost_mask_kernel_compact is 46.5% of GPU time (30 µs/frame). Microbenchmark suggests this kernel can drop to ~35 µs with BLOCK_H=8, BLOCK_W=128, num_warps=2, but end-to-end FPS doesn't change while the CPU dispatch is serial — GPU time already fits inside the sync wait.
• Two follow-ups expected to meaningfully move FPS: (a) a 2-deep frame pipeline that overlaps frame N+1's HtoD with frame N's TRT graph, and (b) moving cv2.findContours off the critical path.