rfdetr-seg: Triton kernel for pre-processing (Jetson Orin NX 8GB) [Includes changes from #2403]#2404
Open
aseembits93 wants to merge 29 commits into
Open
Conversation
Replace the per-frame PIL-bilinear-antialias + to_tensor + normalize chain
in the RF-DETR TRT instance-segmentation model with a single Triton
kernel that resizes, swaps BGR↔RGB, scales by 1/255, and applies
ImageNet normalization — writing straight into the preallocated TRT
input buffer.
Byte-exact port of PIL's separable bilinear-antialias resize
(PRECISION_BITS=22, int32 fixed-point, uint8 quantization between the
horizontal and vertical passes). The horizontal uint8 intermediate
lives in registers.
Correctness
- Preproc max abs error vs PIL: 4.77e-7 (fp32 ULP on the final
/255+normalize step; the uint8 resize result is byte-identical).
- Full coco/val2017 detection parity (rfdetr-seg-nano, conf=0.4):
26,721 / 26,721 matched at IoU>0.5, mean box IoU 1.0000,
|Δscore| 0, 0 class-id disagreements, all matched masks
pixel-identical.
Performance (vehicles_312px.mp4, 538 frames)
- Baseline (PIL path): 76.25 fps
- Triton fast path: 99.83 fps (+31%)
- Preproc microbench (1080p → 312²): 27.0 ms → 2.8 ms per frame (~10×)
Scope
- Gated on: single-image numpy uint8 HWC input, stretch/letterbox/
center-crop/letterbox-reflect resize modes (all collapse to a single
PIL stretch when dataset_version_resize_dimensions is None, verified
via synthetic-package test), no static_crop/grayscale/contrast,
3-channel, scaling_factor in {None, 255}, normalization set.
- Falls back to the existing PIL-based pre_process_network_input
when any precondition fails.
Also adds the benchmark driver
development/stream_interface/rfdetr_nano_seg_trt_workflow.py used to
measure the above numbers.
INFERENCE_MODELS_RFDETR_TRITON_PREPROC_ENABLED (default true). Setting it to false short-circuits _try_fast_preprocess so every call falls back to the PIL reference path — useful for A/B benchmarking and as an escape hatch if the fused kernel is ever implicated in a regression. e2e on vehicles_312px.mp4 (538 frames, rfdetr-seg-nano TRT, mean of 3): ON (default): 98.57 fps OFF (env=false): 76.60 fps Δ: +28.7% / −2.90 ms/frame
aseembits93
requested review from
PawelPeczek-Roboflow,
grzegorz-roboflow,
hansent and
yeldarby
as code owners
aseembits93
requested review from
dkosowski87,
probicheaux and
rafel-roboflow
as code owners
aseembits93
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
It adds the optimized RF-DETR Triton preprocessing path on top of the Triton RLE postprocess branch. End-to-end numbers for this PR therefore represent the cumulative postproc + preproc gain.
The preprocessing fast path is guarded by:
The implementation uses a two-pass Triton resize/normalize flow:
/255+ normalization into the TensorRT input tensorThe low-level kernels live in
triton_preprocess.py. The TensorRT adapter delegates fast-path eligibility, reusable buffer state, warning throttling, and CUDA event handoff toFastPreprocessRuntimeintriton_preprocess_runtime.py. That runtime caches resample tables, pinned host input, device input, scratch buffers, and a ring of output tensors so later pipelined execution can reuse buffers safely. The fast path remains opt-in and falls back to the reference path with aRuntimeWarningwhen the request is outside the supported Triton contract.Type of Change
Testing
Test details: (Make sure Triton is installed in your environment)
Reference command on
main:Candidate command on
opt-preprocess:vehicles_312px.mp4 (538 frames, src 312x176):
vehicles_720p.mp4 (538 frames, src 1280x720):
vehicles_1080p.mp4 (538 frames, src 1920x1080):
mainflags-offTriton replay exercises the production
FastPreprocessRuntimehelper used by the TRT adapter; the harness only materializes captured inputs and compares outputs.Each row uses 100 captured calls and 300 timed replays per implementation. These numbers were rerun after replay was switched to the production
FastPreprocessRuntimehelper.All Triton replay outputs matched captured reference outputs at
atol=1e-6.Results: unit coverage
40 passed; model-level Triton preproc parity1 passedwith the local Orin TRT package override. The default T4 fixture skips cleanly
on this Orin runtime due TensorRT serialized-engine platform mismatch.
Checklist
Additional Context