Skip to content

Improvements to Triton fast path for RF-DETR Instance Segmentation - #2443

Closed
dkosowski87 wants to merge 62 commits into
mainfrom
contrib-pr-2405
Closed

Improvements to Triton fast path for RF-DETR Instance Segmentation#2443
dkosowski87 wants to merge 62 commits into
mainfrom
contrib-pr-2405

Conversation

@dkosowski87

Copy link
Copy Markdown
Contributor

What does this PR do?

This PR is built upon #2405. See commits added by me to check what has been added, plus the description below.

Triton JIT fallback to main path on specific errors during JIT compilation

Context

During live tests RF-DETR TRT Triton pre/post enabled failed on JIT compilation. Examples we hit in the wild:

  • Missing C compiler when Triton compiles driver.c: (Failed to find C compiler. Please specify via CC environment variable)
  • Stale bundled ptxas on new architectures: (ptxas-blackwell fatal: Value 'sm_110a' is not defined for option 'gpu-name')

When that happens, the workflow crashes even though the reference pre/post paths work. Eligibility checks (try_preprocess → None, postproc returning None) only cover unsupported inputs — they do not catch exceptions raised on first kernel compile.

Solution

Add triton_jit_fallback.py and wire it into the Triton fast paths:

  • Preproc triton_preprocess_runtime.py: wrap the Triton kernel launch in try_preprocess. On JIT failure, log an error, set _jit_disabled on the runtime instance, return None so rfdetr_instance_segmentation_trt.pre_process uses the reference path.
  • Postproc common.py: wrap post_process_single_instance_segmentation_result_to_rle_masks_triton in try/except. On JIT failure, log an error, set process-wide _TRITON_POSTPROC_JIT_DISABLED, fall back to the reference RLE implementation.

is_triton_jit_failure() treats these as compile-time / environment failures:

  • PTXASError, OutOfResources (triton.runtime.errors)
  • RuntimeError messages containing: c compiler, ptxas, ptx codegen, triton ptx, out of resource

Unrelated errors (e.g. CUDA OOM) still propagate. After the first JIT failure, Triton is not retried for the rest of the process (instance flag for preproc, module flag for postproc).

Tests

tests/unit_tests/models/rfdetr/test_triton_jit_fallback.py:

  • Marker / exception-type detection (PTXASError, OutOfResources, C compiler, unrelated RuntimeError rejected)
  • One-shot error logging via warn_triton_jit_fallback
  • Preproc: mocked kernel raises C-compiler RuntimeErrorNone, _jit_disabled, no retry on second call
  • Postproc: mocked Triton path raises → reference result, _TRITON_POSTPROC_JIT_DISABLED, no retry on second call

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 and others added 27 commits June 10, 2026 06:30
- Introduced `triton_jit_fallback.py` to detect Triton JIT compilation failures and log warnings.
- Updated `common.py` and `triton_preprocess_runtime.py` to utilize the new fallback mechanisms, ensuring graceful degradation to reference paths on failure.
- Added unit tests for Triton JIT failure detection and fallback behavior to ensure robustness.
- Added handling for OutOfResources exceptions in `triton_jit_fallback.py`.
- Updated the `is_triton_jit_failure` function to include OutOfResources in its checks.
- Introduced new unit tests to verify detection of OutOfResources errors and related messages.
- Added 'opsx/' and '.cursor/' to the .gitignore file to prevent tracking of these directories.
- Ensured proper formatting by adding a newline at the end of the file.
- Updated exception handling to define `_TRITON_JIT_EXCEPTION_TYPES` as a tuple of `_PTXASError` and `_OutOfResources`.
- Ensured that the exception types are set to an empty tuple in case of an ImportError, improving clarity and robustness of the error handling logic.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants