From bb1abe68722ca14de36438c98802a9d2347d56be Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Sat, 9 May 2026 00:30:01 +0000 Subject: [PATCH 01/24] perf(rfdetr-seg): fused Triton preproc kernel for TRT path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../rfdetr_nano_seg_trt_workflow.py | 116 ++++++ .../rfdetr_instance_segmentation_trt.py | 223 ++++++++++ .../models/rfdetr/triton_preprocess.py | 383 ++++++++++++++++++ 3 files changed, 722 insertions(+) create mode 100644 development/stream_interface/rfdetr_nano_seg_trt_workflow.py create mode 100644 inference_models/inference_models/models/rfdetr/triton_preprocess.py diff --git a/development/stream_interface/rfdetr_nano_seg_trt_workflow.py b/development/stream_interface/rfdetr_nano_seg_trt_workflow.py new file mode 100644 index 0000000000..9c213a8639 --- /dev/null +++ b/development/stream_interface/rfdetr_nano_seg_trt_workflow.py @@ -0,0 +1,116 @@ +"""Minimal benchmark: RF-DETR instance segmentation through inference-models, +run via InferencePipeline on a single video source. + +Workflow has exactly one block — the segmentation model. No annotators, no +buffer strategies, no rate limiting. + +The `--backend` flag (trt | onnx | torch) is parsed before importing +`inference` and pins the auto-loader by setting +`DISABLED_INFERENCE_MODELS_BACKENDS` to every backend except the chosen one, +so the benchmark numbers correspond unambiguously to a single execution path. + +Defaults: rfdetr-seg-nano @ confidence 0.4 on the native TRT backend. +""" +import argparse +import os + +_ALL_BACKENDS = { + "torch", + "torch-script", + "onnx", + "trt", + "hugging-face", + "ultralytics", + "mediapipe", + "custom", +} +def _select_backend_from_argv() -> str: + pre = argparse.ArgumentParser(add_help=False) + pre.add_argument("--backend", choices=("trt", "onnx", "torch"), default="trt") + args, _ = pre.parse_known_args() + return args.backend + + +_BACKEND = _select_backend_from_argv() +os.environ.setdefault( + "ONNXRUNTIME_EXECUTION_PROVIDERS", + "[TensorrtExecutionProvider,CUDAExecutionProvider,CPUExecutionProvider]", +) +os.environ["DISABLED_INFERENCE_MODELS_BACKENDS"] = ",".join( + sorted(_ALL_BACKENDS - {_BACKEND}) +) + +from time import perf_counter + +from inference import InferencePipeline + + +def build_workflow(model_id: str, confidence: float) -> dict: + return { + "version": "1.0", + "inputs": [{"type": "WorkflowImage", "name": "image"}], + "steps": [ + { + "type": "roboflow_core/roboflow_instance_segmentation_model@v3", + "name": "segmentation", + "images": "$inputs.image", + "model_id": model_id, + "confidence_mode": "custom", + "custom_confidence": confidence, + }, + ], + "outputs": [ + { + "type": "JsonField", + "name": "predictions", + "selector": "$steps.segmentation.predictions", + }, + ], + } + +FRAME_COUNT = 0 +START_TIME = None +PROGRESS_EVERY = 50 + + +def sink(predictions, _video_frames) -> None: + global FRAME_COUNT, START_TIME + del _video_frames + if not isinstance(predictions, list): + predictions = [predictions] + FRAME_COUNT += sum(p is not None for p in predictions) + if START_TIME is None: + START_TIME = perf_counter() + if FRAME_COUNT % PROGRESS_EVERY == 0: + fps = FRAME_COUNT / (perf_counter() - START_TIME) + print(f"[progress] frames={FRAME_COUNT} fps={fps:.2f}", flush=True) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--video_reference", required=True) + parser.add_argument("--model_id", default="rfdetr-seg-nano") + parser.add_argument("--confidence", type=float, default=0.4) + parser.add_argument( + "--backend", + choices=("trt", "onnx", "torch"), + default="trt", + help="inference-models backend (consumed pre-import via env var).", + ) + args = parser.parse_args() + + pipeline = InferencePipeline.init_with_workflow( + video_reference=args.video_reference, + workflow_specification=build_workflow(args.model_id, args.confidence), + on_prediction=sink, + ) + pipeline.start() + pipeline.join() + + elapsed = perf_counter() - START_TIME if START_TIME else 0.0 + fps = FRAME_COUNT / elapsed if elapsed > 0 else 0.0 + print(f"frames={FRAME_COUNT} elapsed={elapsed:.2f}s fps={fps:.2f}") + + +if __name__ == "__main__": + main() diff --git a/inference_models/inference_models/models/rfdetr/rfdetr_instance_segmentation_trt.py b/inference_models/inference_models/models/rfdetr/rfdetr_instance_segmentation_trt.py index e8196f8eae..afa70f9c3d 100644 --- a/inference_models/inference_models/models/rfdetr/rfdetr_instance_segmentation_trt.py +++ b/inference_models/inference_models/models/rfdetr/rfdetr_instance_segmentation_trt.py @@ -26,10 +26,13 @@ use_primary_cuda_context, ) from inference_models.models.common.model_packages import get_model_package_contents +from inference_models.entities import ImageDimensions from inference_models.models.common.roboflow.model_packages import ( + ColorMode, InferenceConfig, PreProcessingMetadata, ResizeMode, + StaticCropOffset, TRTConfig, parse_class_names_file, parse_inference_config, @@ -52,6 +55,17 @@ post_process_instance_segmentation_results_to_rle_masks, ) from inference_models.models.rfdetr.pre_processing import pre_process_network_input + +try: + from inference_models.models.rfdetr.triton_preprocess import ( + TRITON_AVAILABLE as _TRITON_AVAILABLE, + build_resample_tables, + triton_preprocess_rfdetr_stretch, + ) +except ImportError: + _TRITON_AVAILABLE = False + build_resample_tables = None + triton_preprocess_rfdetr_stretch = None from inference_models.weights_providers.entities import RecommendedParameters try: @@ -85,6 +99,84 @@ ) from import_error +class _FastPathState: + """Per-(src_shape, target_shape) cache of GPU buffers + resample tables + that the Triton fast path reuses across frames.""" + + __slots__ = ( + "src_h", + "src_w", + "target_h", + "target_w", + "pinned_host", + "src_gpu", + "out_buffer", + "tables", + ) + + def __init__( + self, + src_h: int, + src_w: int, + target_h: int, + target_w: int, + pinned_host: torch.Tensor, + src_gpu: torch.Tensor, + out_buffer: torch.Tensor, + tables, + ) -> None: + self.src_h = src_h + self.src_w = src_w + self.target_h = target_h + self.target_w = target_w + self.pinned_host = pinned_host + self.src_gpu = src_gpu + self.out_buffer = out_buffer + self.tables = tables + + @classmethod + def build( + cls, + src_h: int, + src_w: int, + target_h: int, + target_w: int, + device: torch.device, + ) -> "_FastPathState": + pinned_host = torch.empty((src_h, src_w, 3), dtype=torch.uint8, pin_memory=True) + src_gpu = torch.empty((src_h, src_w, 3), dtype=torch.uint8, device=device) + out_buffer = torch.empty( + (1, 3, target_h, target_w), dtype=torch.float32, device=device + ) + tables = build_resample_tables( + src_h=src_h, + src_w=src_w, + target_h=target_h, + target_w=target_w, + device=device, + ) + return cls( + src_h=src_h, + src_w=src_w, + target_h=target_h, + target_w=target_w, + pinned_host=pinned_host, + src_gpu=src_gpu, + out_buffer=out_buffer, + tables=tables, + ) + + def is_stale( + self, src_h: int, src_w: int, target_h: int, target_w: int + ) -> bool: + return ( + self.src_h != src_h + or self.src_w != src_w + or self.target_h != target_h + or self.target_w != target_w + ) + + class RFDetrForInstanceSegmentationTRT( InstanceSegmentationModel[ torch.Tensor, @@ -220,6 +312,7 @@ def __init__( self._inference_stream = torch.cuda.Stream(device=self._device) self._thread_local_storage = threading.local() self.recommended_parameters = recommended_parameters + self._fast_path_state: Optional[_FastPathState] = None @property def class_names(self) -> List[str]: @@ -237,6 +330,14 @@ def pre_process( pre_processing_overrides: Optional[PreProcessingOverrides] = None, **kwargs, ) -> Tuple[torch.Tensor, List[PreProcessingMetadata]]: + fast = self._try_fast_preprocess( + images=images, + input_color_format=input_color_format, + image_size=image_size, + pre_processing_overrides=pre_processing_overrides, + ) + if fast is not None: + return fast with torch.cuda.stream(self._pre_process_stream): pre_processed_images, pre_processing_meta = pre_process_network_input( images=images, @@ -250,6 +351,128 @@ def pre_process( self._pre_process_stream.synchronize() return pre_processed_images, pre_processing_meta + def _try_fast_preprocess( + self, + images, + input_color_format, + image_size, + pre_processing_overrides, + ) -> Optional[Tuple[torch.Tensor, List[PreProcessingMetadata]]]: + if not _TRITON_AVAILABLE: + return None + if image_size is not None: + return None + # pre_processing_overrides can only *disable* transforms; it has no + # "enable" knob. The fast path never applies static_crop / grayscale / + # contrast regardless, so the override flags are irrelevant — we just + # gate on whether the image_pre_processing config itself asks for them. + ipp = self._inference_config.image_pre_processing + if ( + (ipp.static_crop is not None and ipp.static_crop.enabled) + or (ipp.contrast is not None and ipp.contrast.enabled) + or (ipp.grayscale is not None and ipp.grayscale.enabled) + ): + return None + + ni = self._inference_config.network_input + if ni.dataset_version_resize_dimensions is not None: + return None + if ni.input_channels != 3: + return None + if ni.scaling_factor not in (None, 255): + return None + if ni.normalization is None: + return None + # When dataset_version_resize_dimensions is None, the prod path collapses + # non-stretch resize modes to a single PIL stretch as well + # (pre_processing.py:_needs_two_step_resize), so we accept all modes here. + if ni.resize_mode not in ( + ResizeMode.STRETCH_TO, + ResizeMode.LETTERBOX, + ResizeMode.CENTER_CROP, + ResizeMode.LETTERBOX_REFLECT_EDGES, + ): + return None + + if isinstance(images, list): + if len(images) != 1: + return None + candidate = images[0] + else: + candidate = images + if not isinstance(candidate, np.ndarray): + return None + if ( + candidate.dtype != np.uint8 + or candidate.ndim != 3 + or candidate.shape[2] != 3 + ): + return None + + caller_mode = ( + ColorMode(input_color_format) + if input_color_format is not None + else ColorMode.BGR + ) + swap_rb = caller_mode != ni.color_mode + + means, stds = ni.normalization + means_t = (float(means[0]), float(means[1]), float(means[2])) + stds_t = (float(stds[0]), float(stds[1]), float(stds[2])) + target_h = ni.training_input_size.height + target_w = ni.training_input_size.width + orig_h, orig_w = int(candidate.shape[0]), int(candidate.shape[1]) + + state = self._fast_path_state + if state is None or state.is_stale( + src_h=orig_h, + src_w=orig_w, + target_h=target_h, + target_w=target_w, + ): + state = _FastPathState.build( + src_h=orig_h, + src_w=orig_w, + target_h=target_h, + target_w=target_w, + device=self._device, + ) + self._fast_path_state = state + + pinned_np = state.pinned_host.numpy() + np.copyto(pinned_np, candidate, casting="no") + + with torch.cuda.stream(self._pre_process_stream): + state.src_gpu.copy_(state.pinned_host, non_blocking=True) + triton_preprocess_rfdetr_stretch( + src=state.src_gpu, + tables=state.tables, + target_h=target_h, + target_w=target_w, + means=means_t, + stds=stds_t, + swap_rb=swap_rb, + out=state.out_buffer, + ) + state.out_buffer.record_stream(self._pre_process_stream) + self._pre_process_stream.synchronize() + + meta = PreProcessingMetadata( + pad_left=0, + pad_top=0, + pad_right=0, + pad_bottom=0, + original_size=ImageDimensions(width=orig_w, height=orig_h), + size_after_pre_processing=ImageDimensions(width=orig_w, height=orig_h), + inference_size=ImageDimensions(width=target_w, height=target_h), + scale_width=target_w / orig_w, + scale_height=target_h / orig_h, + static_crop_offset=StaticCropOffset( + offset_x=0, offset_y=0, crop_width=orig_w, crop_height=orig_h + ), + ) + return state.out_buffer, [meta] + def forward( self, pre_processed_images: torch.Tensor, diff --git a/inference_models/inference_models/models/rfdetr/triton_preprocess.py b/inference_models/inference_models/models/rfdetr/triton_preprocess.py new file mode 100644 index 0000000000..0b9d7eaa8e --- /dev/null +++ b/inference_models/inference_models/models/rfdetr/triton_preprocess.py @@ -0,0 +1,383 @@ +"""Fused Triton preprocessing kernel for RF-DETR. + +Byte-exact port of PIL's separable bilinear-antialias resize (the algorithm +torchvision's `TF.resize(pil, ..., antialias=True)` uses on PIL inputs), with +the subsequent `/255` + ImageNet normalize fused into the same pass. + +PIL's scheme (src/libImaging/Resample.c): + + PRECISION_BITS = 22 + scale = in_size / out_size + filterscale = max(1.0, scale) + support = 1.0 * filterscale # triangle radius = 1 + ksize = ceil(support) * 2 + 1 + center(o) = (o + 0.5) * scale + xmin(o) = int(center - support + 0.5) clipped to [0, in] + xmax(o) = int(center + support + 0.5) clipped to [0, in] + w_f(o, k) = triangle((k + xmin - center + 0.5) / filterscale) + w_f normalised to sum to 1 per output pixel + w_i(o, k) = round(w_f(o, k) * (1 << PRECISION_BITS)) int32 + out(o) = clamp((Σ w_i(o, k) * src_u8) + (1 << (PRECISION_BITS-1)) >> PRECISION_BITS, 0, 255) + +Single fused kernel: the horizontal uint8 intermediate lives in registers +rather than a DRAM scratch buffer. For each output tile we loop over +KSIZE_Y source rows; for each contributing source row we recompute the +horizontal convolution (int32 fixed-point, uint8 quantize) on the fly, +multiply by the vertical weight, and accumulate. Final: uint8 quantize, +BGR↔RGB swap, /255, ImageNet normalize, fp32 CHW store. + +A separable two-pass variant (horizontal then vertical, via a DRAM uint8 +intermediate) is ~0.4 fps faster end-to-end on the 312² RF-DETR workload +because it avoids redoing KSIZE_X MACs per output row. We picked the fused +version for simplicity (no intermediate buffer, one launch, one piece of +math). +""" + +from __future__ import annotations + +import math +from typing import Optional, Tuple + +import numpy as np +import torch + +try: + import triton + import triton.language as tl + + TRITON_AVAILABLE = True +except ImportError: # pragma: no cover + triton = None + tl = None + TRITON_AVAILABLE = False + + +PRECISION_BITS = 22 + + +def _bilinear_antialias_weights_1d_int( + in_size: int, out_size: int +) -> Tuple[np.ndarray, np.ndarray, int]: + """PIL's precompute_coeffs, int32 fixed-point form.""" + scale = in_size / out_size + filterscale = max(1.0, scale) + support = filterscale + ksize = int(math.ceil(support)) * 2 + 1 + + starts = np.zeros(out_size, dtype=np.int32) + weights_fp = np.zeros((out_size, ksize), dtype=np.float64) + inv_fs = 1.0 / filterscale + + for o in range(out_size): + center = (o + 0.5) * scale + xmin = int(center - support + 0.5) + if xmin < 0: + xmin = 0 + xmax = int(center + support + 0.5) + if xmax > in_size: + xmax = in_size + actual = xmax - xmin + starts[o] = xmin + total = 0.0 + for k in range(actual): + t = (k + xmin - center + 0.5) * inv_fs + t_abs = -t if t < 0.0 else t + w = 1.0 - t_abs if t_abs < 1.0 else 0.0 + weights_fp[o, k] = w + total += w + if total != 0.0: + weights_fp[o, :actual] /= total + + weights_int = np.rint(weights_fp * (1 << PRECISION_BITS)).astype(np.int32) + return starts, weights_int, ksize + + +if TRITON_AVAILABLE: + + _HALF = 1 << (PRECISION_BITS - 1) + + @triton.jit + def _fused_resize_normalize_kernel( + src_ptr, + dst_ptr, + ymin_ptr, + xmin_ptr, + wy_ptr, + wx_ptr, + src_h, + src_w, + src_stride_h, + src_stride_w, + dst_stride_c, + dst_stride_h, + target_h, + target_w, + inv_std_255_r, + inv_std_255_g, + inv_std_255_b, + offset_r, + offset_g, + offset_b, + CH_R: tl.constexpr, + CH_G: tl.constexpr, + CH_B: tl.constexpr, + KSIZE_Y: tl.constexpr, + KSIZE_X: tl.constexpr, + PRECISION_BITS_C: tl.constexpr, + HALF_C: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_W: tl.constexpr, + ): + """One kernel per (tile_y, tile_x) over target image. + + In : src uint8 HWC (src_h, src_w, 3), source color order. + Out: dst fp32 CHW (1, 3, target_h, target_w), network color order, + (pixel/255 - mean)/std. + """ + pid_y = tl.program_id(0) + pid_x = tl.program_id(1) + + offs_y = pid_y * BLOCK_H + tl.arange(0, BLOCK_H) + offs_x = pid_x * BLOCK_W + tl.arange(0, BLOCK_W) + mask_y = offs_y < target_h + mask_x = offs_x < target_w + mask_out = mask_y[:, None] & mask_x[None, :] + + ymin = tl.load(ymin_ptr + offs_y, mask=mask_y, other=0) + xmin = tl.load(xmin_ptr + offs_x, mask=mask_x, other=0) + + # Vertical pass accumulators (int32 fixed-point) for 3 channels. + vacc_0 = tl.zeros((BLOCK_H, BLOCK_W), dtype=tl.int32) + vacc_1 = tl.zeros((BLOCK_H, BLOCK_W), dtype=tl.int32) + vacc_2 = tl.zeros((BLOCK_H, BLOCK_W), dtype=tl.int32) + + for ky in tl.static_range(KSIZE_Y): + # Source row contributing to each output row in this tile. + sy = ymin + ky + sy_c = tl.maximum(tl.minimum(sy, src_h - 1), 0) + wy = tl.load(wy_ptr + offs_y * KSIZE_Y + ky, mask=mask_y, other=0) + + # Horizontal pass for (output_rows_in_tile, output_cols_in_tile): + # for each source column in the kernel, gather src[sy_c, sx_c, :] + # and accumulate with wx[output_col, kx]. + hacc_0 = tl.zeros((BLOCK_H, BLOCK_W), dtype=tl.int32) + hacc_1 = tl.zeros((BLOCK_H, BLOCK_W), dtype=tl.int32) + hacc_2 = tl.zeros((BLOCK_H, BLOCK_W), dtype=tl.int32) + + for kx in tl.static_range(KSIZE_X): + sx = xmin + kx + sx_c = tl.maximum(tl.minimum(sx, src_w - 1), 0) + wx = tl.load(wx_ptr + offs_x * KSIZE_X + kx, mask=mask_x, other=0) + base = sy_c[:, None] * src_stride_h + sx_c[None, :] * src_stride_w + p0 = tl.load(src_ptr + base + 0, mask=mask_out, other=0).to(tl.int32) + p1 = tl.load(src_ptr + base + 1, mask=mask_out, other=0).to(tl.int32) + p2 = tl.load(src_ptr + base + 2, mask=mask_out, other=0).to(tl.int32) + wx_2d = wx[None, :] + hacc_0 += p0 * wx_2d + hacc_1 += p1 * wx_2d + hacc_2 += p2 * wx_2d + + # Horizontal uint8 quantization (byte-exact to PIL's intermediate). + hacc_0 = (hacc_0 + HALF_C) >> PRECISION_BITS_C + hacc_1 = (hacc_1 + HALF_C) >> PRECISION_BITS_C + hacc_2 = (hacc_2 + HALF_C) >> PRECISION_BITS_C + hacc_0 = tl.minimum(tl.maximum(hacc_0, 0), 255) + hacc_1 = tl.minimum(tl.maximum(hacc_1, 0), 255) + hacc_2 = tl.minimum(tl.maximum(hacc_2, 0), 255) + + wy_2d = wy[:, None] + vacc_0 += hacc_0 * wy_2d + vacc_1 += hacc_1 * wy_2d + vacc_2 += hacc_2 * wy_2d + + # Vertical uint8 quantization. + q_0 = (vacc_0 + HALF_C) >> PRECISION_BITS_C + q_1 = (vacc_1 + HALF_C) >> PRECISION_BITS_C + q_2 = (vacc_2 + HALF_C) >> PRECISION_BITS_C + q_0 = tl.minimum(tl.maximum(q_0, 0), 255) + q_1 = tl.minimum(tl.maximum(q_1, 0), 255) + q_2 = tl.minimum(tl.maximum(q_2, 0), 255) + + # Source-to-output channel remap (triton requires constexpr branches). + if CH_R == 0: + q_r = q_0 + elif CH_R == 1: + q_r = q_1 + else: + q_r = q_2 + if CH_G == 0: + q_g = q_0 + elif CH_G == 1: + q_g = q_1 + else: + q_g = q_2 + if CH_B == 0: + q_b = q_0 + elif CH_B == 1: + q_b = q_1 + else: + q_b = q_2 + + # (pixel/255 - mean)/std == pixel * (1/(255*std)) + (-mean/std) + out_r = q_r.to(tl.float32) * inv_std_255_r + offset_r + out_g = q_g.to(tl.float32) * inv_std_255_g + offset_g + out_b = q_b.to(tl.float32) * inv_std_255_b + offset_b + + out_row = offs_y[:, None] * dst_stride_h + offs_x[None, :] + tl.store(dst_ptr + 0 * dst_stride_c + out_row, out_r, mask=mask_out) + tl.store(dst_ptr + 1 * dst_stride_c + out_row, out_g, mask=mask_out) + tl.store(dst_ptr + 2 * dst_stride_c + out_row, out_b, mask=mask_out) + + +class _ResampleTables: + """Cache of per-axis PIL-int32 weight tables for one (src, dst) pair.""" + + __slots__ = ( + "ymin_gpu", + "xmin_gpu", + "wy_gpu", + "wx_gpu", + "ksize_y", + "ksize_x", + ) + + def __init__( + self, + ymin_gpu: torch.Tensor, + xmin_gpu: torch.Tensor, + wy_gpu: torch.Tensor, + wx_gpu: torch.Tensor, + ksize_y: int, + ksize_x: int, + ) -> None: + self.ymin_gpu = ymin_gpu + self.xmin_gpu = xmin_gpu + self.wy_gpu = wy_gpu + self.wx_gpu = wx_gpu + self.ksize_y = ksize_y + self.ksize_x = ksize_x + + +def build_resample_tables( + src_h: int, + src_w: int, + target_h: int, + target_w: int, + device: torch.device, +) -> _ResampleTables: + ymin, wy, ksize_y = _bilinear_antialias_weights_1d_int(src_h, target_h) + xmin, wx, ksize_x = _bilinear_antialias_weights_1d_int(src_w, target_w) + return _ResampleTables( + ymin_gpu=torch.from_numpy(ymin).to(device=device, non_blocking=True), + xmin_gpu=torch.from_numpy(xmin).to(device=device, non_blocking=True), + wy_gpu=torch.from_numpy(wy.ravel()).to(device=device, non_blocking=True), + wx_gpu=torch.from_numpy(wx.ravel()).to(device=device, non_blocking=True), + ksize_y=ksize_y, + ksize_x=ksize_x, + ) + + +def triton_preprocess_rfdetr_stretch( + src: torch.Tensor, + tables: _ResampleTables, + target_h: int, + target_w: int, + means: Tuple[float, float, float] = (0.485, 0.456, 0.406), + stds: Tuple[float, float, float] = (0.229, 0.224, 0.225), + swap_rb: bool = True, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Fused PIL-exact resize + color swap + normalize. + + Args: + src: uint8 CUDA tensor, shape (H, W, 3), HWC layout. + tables: precomputed int32 resample tables from `build_resample_tables`. + target_h, target_w: output spatial dims. + means, stds: normalization in output channel order (R, G, B for + network_input.color_mode == 'rgb'). + swap_rb: if True, source channel 0 → output B (BGR input, RGB network). + out: optional preallocated fp32 (1, 3, H, W) CUDA tensor. + + Returns: + fp32 (1, 3, target_h, target_w) on the same device as `src`. + """ + if not TRITON_AVAILABLE: + raise RuntimeError("triton is not installed") + if not src.is_cuda: + raise ValueError(f"expected CUDA src tensor, got device={src.device}") + if src.dtype != torch.uint8: + raise ValueError(f"expected uint8 src, got {src.dtype}") + if src.ndim != 3 or src.shape[2] != 3: + raise ValueError(f"expected HWC 3-channel, got shape={tuple(src.shape)}") + + src = src.contiguous() + src_h, src_w = int(src.shape[0]), int(src.shape[1]) + src_stride_h = int(src.stride(0)) + src_stride_w = int(src.stride(1)) + + if out is None: + out = torch.empty( + (1, 3, target_h, target_w), dtype=torch.float32, device=src.device + ) + else: + if tuple(out.shape) != (1, 3, target_h, target_w): + raise ValueError( + f"out has shape {tuple(out.shape)}, expected " + f"(1, 3, {target_h}, {target_w})" + ) + if out.dtype != torch.float32 or not out.is_cuda: + raise ValueError("out must be fp32 CUDA tensor") + + dst_stride_c = target_h * target_w + dst_stride_h = target_w + + if swap_rb: + ch_r, ch_g, ch_b = 2, 1, 0 + else: + ch_r, ch_g, ch_b = 0, 1, 2 + + inv_std_255_r = 1.0 / (255.0 * stds[0]) + inv_std_255_g = 1.0 / (255.0 * stds[1]) + inv_std_255_b = 1.0 / (255.0 * stds[2]) + offset_r = -means[0] / stds[0] + offset_g = -means[1] / stds[1] + offset_b = -means[2] / stds[2] + + BLOCK_H = 16 + BLOCK_W = 16 + grid = ( + (target_h + BLOCK_H - 1) // BLOCK_H, + (target_w + BLOCK_W - 1) // BLOCK_W, + ) + _fused_resize_normalize_kernel[grid]( + src, + out, + tables.ymin_gpu, + tables.xmin_gpu, + tables.wy_gpu, + tables.wx_gpu, + src_h, + src_w, + src_stride_h, + src_stride_w, + dst_stride_c, + dst_stride_h, + target_h, + target_w, + float(inv_std_255_r), + float(inv_std_255_g), + float(inv_std_255_b), + float(offset_r), + float(offset_g), + float(offset_b), + CH_R=ch_r, + CH_G=ch_g, + CH_B=ch_b, + KSIZE_Y=tables.ksize_y, + KSIZE_X=tables.ksize_x, + PRECISION_BITS_C=PRECISION_BITS, + HALF_C=_HALF, + BLOCK_H=BLOCK_H, + BLOCK_W=BLOCK_W, + ) + return out From 3e4892a97ad0caf5adfd43ef9d3cb2c1a8cb9062 Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Sat, 9 May 2026 02:46:38 +0000 Subject: [PATCH 02/24] =?UTF-8?q?perf(rfdetr):=20widen=20fast=20path=20?= =?UTF-8?q?=E2=80=94=20single=20integration=20point=20+=20torch+batched=20?= =?UTF-8?q?input?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move the Triton fast-path gate from RFDetrForInstanceSegmentationTRT into pre_process_network_input so all six RFDetr classes (seg×{TRT,ONNX,Torch} and detect×{TRT,ONNX,Torch}) can hit it, and widen the predicate to accept torch uint8 HWC tensors on any device plus batched inputs (list[ndarray], list[Tensor], 4D ndarray/Tensor — the outer function already unbinds those to lists before the per-item check). Color-swap parity fix: the PIL path does `image[:, :, ::-1]` whenever `input_color_mode != network_input.color_mode`, which is True for an unspecified caller (None). The old fast-path treated None as BGR and skipped the swap when the network was also BGR — byte-identical to PIL for packaged seg models but diverged from PIL on og-rfdetr-base (ColorMode.BGR network with None caller). Align the kernel swap condition with PIL's. Integration coverage (144 tests, CUDA 13): baseline: 4 tests hit fast path, 6 / 160 pre_process calls widened : 35 tests hit fast path, 43 / 166 pre_process calls pass rate unchanged: 144 / 144. Remaining ~100 tests miss on predicate categories that require kernel extensions (static_crop, contrast, dataset_version_resize) and are tracked as follow-up work. --- .../models/rfdetr/pre_processing.py | 184 ++++++++++++++- .../rfdetr_instance_segmentation_trt.py | 223 ------------------ 2 files changed, 183 insertions(+), 224 deletions(-) diff --git a/inference_models/inference_models/models/rfdetr/pre_processing.py b/inference_models/inference_models/models/rfdetr/pre_processing.py index f0fbe586db..8a1120c771 100644 --- a/inference_models/inference_models/models/rfdetr/pre_processing.py +++ b/inference_models/inference_models/models/rfdetr/pre_processing.py @@ -9,9 +9,15 @@ torch.Tensor inputs (advanced caller, float CHW [0, 1]): tensor F.resize → F.normalize + +Fused Triton fast path: for the common case (single-stage resize, no +image-level transforms, normalization set, uint8 HWC input) we invoke a +single PIL-exact Triton kernel that writes the normalized fp32 tensor +directly to `target_device`. See +`inference_models.models.rfdetr.triton_preprocess`. """ -from typing import List, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import numpy as np import torch @@ -40,6 +46,34 @@ pre_process_numpy_image, ) +try: + from inference_models.models.rfdetr.triton_preprocess import ( + TRITON_AVAILABLE as _TRITON_AVAILABLE, + build_resample_tables, + triton_preprocess_rfdetr_stretch, + ) +except ImportError: + _TRITON_AVAILABLE = False + build_resample_tables = None + triton_preprocess_rfdetr_stretch = None + +# Resample tables are cheap (a few KB of int32 weights) but rebuilding them +# is a per-call Python loop we can skip. Key: (device_str, src_h, src_w, th, tw). +_RESAMPLE_TABLES_CACHE: Dict[Tuple[str, int, int, int, int], "object"] = {} + + +def _get_resample_tables( + device: torch.device, src_h: int, src_w: int, th: int, tw: int +): + key = (str(device), src_h, src_w, th, tw) + t = _RESAMPLE_TABLES_CACHE.get(key) + if t is None: + t = build_resample_tables( + src_h=src_h, src_w=src_w, target_h=th, target_w=tw, device=device + ) + _RESAMPLE_TABLES_CACHE[key] = t + return t + def pre_process_network_input( images: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]], @@ -92,6 +126,21 @@ def pre_process_network_input( else: image_list = [images] + if _fast_path_eligible( + image_list=image_list, + image_pre_processing=image_pre_processing, + network_input=network_input, + target_device=target_device, + pre_processing_overrides=pre_processing_overrides, + ): + return _fast_path_preprocess( + image_list=image_list, + network_input=network_input, + target_size=target_size, + target_device=target_device, + input_color_mode=input_color_mode, + ) + tensors: List[torch.Tensor] = [] metadata: List[PreProcessingMetadata] = [] for img in image_list: @@ -127,6 +176,139 @@ def pre_process_network_input( return batch, metadata +def _fast_path_eligible( + image_list: List[Union[np.ndarray, torch.Tensor]], + image_pre_processing: ImagePreProcessing, + network_input: NetworkInputDefinition, + target_device: torch.device, + pre_processing_overrides: Optional[PreProcessingOverrides], +) -> bool: + """True when every item in `image_list` can go through the fused Triton + kernel. Predicate is intentionally conservative — any miss → PIL path.""" + if not _TRITON_AVAILABLE: + return False + # Kernel runs on CUDA. + if target_device.type != "cuda": + return False + # Image-level transforms (static_crop / grayscale / contrast) aren't in + # the kernel — predicate rejects if any is enabled in config. + ipp = image_pre_processing + if ( + (ipp.static_crop is not None and ipp.static_crop.enabled) + or (ipp.contrast is not None and ipp.contrast.enabled) + or (ipp.grayscale is not None and ipp.grayscale.enabled) + ): + return False + # Two-stage dataset-version resize isn't in the kernel. + ni = network_input + if ni.dataset_version_resize_dimensions is not None: + return False + if ni.input_channels != 3: + return False + if ni.scaling_factor not in (None, 255): + return False + if ni.normalization is None: + return False + if ni.resize_mode not in ( + ResizeMode.STRETCH_TO, + ResizeMode.LETTERBOX, + ResizeMode.CENTER_CROP, + ResizeMode.LETTERBOX_REFLECT_EDGES, + ): + return False + # When dataset_version_resize_dimensions is None, the prod PIL path + # collapses every resize_mode to a single stretch to training_input_size; + # the kernel does the same. See _needs_two_step_resize. + if not image_list: + return False + # `pre_process_network_input` already unbinds 4D inputs (batch tensors + # and 4D numpy arrays) into a list of 3D items before calling us, so + # the per-item check below only needs to handle 3D. + for img in image_list: + if isinstance(img, np.ndarray): + if img.dtype != np.uint8 or img.ndim != 3 or img.shape[2] != 3: + return False + elif isinstance(img, torch.Tensor): + # Only uint8 HWC. Float tensors keep the existing tensor branch + # (F.resize bilinear *without* PIL's antialias — a caller- + # accepted divergence we don't silently change). + if img.dtype != torch.uint8: + return False + if img.ndim != 3 or img.shape[-1] != 3: + return False + else: + return False + return True + + +def _as_hwc_uint8_cuda( + img: Union[np.ndarray, torch.Tensor], device: torch.device +) -> torch.Tensor: + """Return a contiguous (H, W, 3) uint8 CUDA tensor, copying if needed.""" + if isinstance(img, torch.Tensor): + if img.device != device: + img = img.to(device=device, non_blocking=True) + return img.contiguous() + t = torch.from_numpy(np.ascontiguousarray(img)) + return t.to(device=device, non_blocking=True) + + +def _fast_path_preprocess( + image_list: List[Union[np.ndarray, torch.Tensor]], + network_input: NetworkInputDefinition, + target_size: ImageDimensions, + target_device: torch.device, + input_color_mode: Optional[ColorMode], +) -> Tuple[torch.Tensor, List[PreProcessingMetadata]]: + """Per-image Triton launch + stack. Assumes `_fast_path_eligible` passed.""" + means, stds = network_input.normalization + means_t = (float(means[0]), float(means[1]), float(means[2])) + stds_t = (float(stds[0]), float(stds[1]), float(stds[2])) + th, tw = target_size.height, target_size.width + + # Match _pre_process_numpy's swap semantics: the PIL path does + # `image[:, :, ::-1]` whenever `input_color_mode != network_input.color_mode`. + # That comparison returns True when input_color_mode is None (unspecified), + # so the swap runs even on unspecified inputs — we replicate that exactly. + swap_rb = input_color_mode != network_input.color_mode + + outs: List[torch.Tensor] = [] + metas: List[PreProcessingMetadata] = [] + for img in image_list: + src_gpu = _as_hwc_uint8_cuda(img, target_device) + sh, sw = int(src_gpu.shape[0]), int(src_gpu.shape[1]) + tables = _get_resample_tables(target_device, sh, sw, th, tw) + out = triton_preprocess_rfdetr_stretch( + src=src_gpu, + tables=tables, + target_h=th, + target_w=tw, + means=means_t, + stds=stds_t, + swap_rb=swap_rb, + ) + outs.append(out[0]) # drop leading batch dim to match stack below + metas.append( + PreProcessingMetadata( + pad_left=0, + pad_top=0, + pad_right=0, + pad_bottom=0, + original_size=ImageDimensions(width=sw, height=sh), + size_after_pre_processing=ImageDimensions(width=sw, height=sh), + inference_size=ImageDimensions(width=tw, height=th), + scale_width=tw / sw, + scale_height=th / sh, + static_crop_offset=StaticCropOffset( + offset_x=0, offset_y=0, crop_width=sw, crop_height=sh + ), + ) + ) + + batch = torch.stack(outs).contiguous() + return batch, metas + + def _pre_process_numpy( image: np.ndarray, image_pre_processing: ImagePreProcessing, diff --git a/inference_models/inference_models/models/rfdetr/rfdetr_instance_segmentation_trt.py b/inference_models/inference_models/models/rfdetr/rfdetr_instance_segmentation_trt.py index afa70f9c3d..e8196f8eae 100644 --- a/inference_models/inference_models/models/rfdetr/rfdetr_instance_segmentation_trt.py +++ b/inference_models/inference_models/models/rfdetr/rfdetr_instance_segmentation_trt.py @@ -26,13 +26,10 @@ use_primary_cuda_context, ) from inference_models.models.common.model_packages import get_model_package_contents -from inference_models.entities import ImageDimensions from inference_models.models.common.roboflow.model_packages import ( - ColorMode, InferenceConfig, PreProcessingMetadata, ResizeMode, - StaticCropOffset, TRTConfig, parse_class_names_file, parse_inference_config, @@ -55,17 +52,6 @@ post_process_instance_segmentation_results_to_rle_masks, ) from inference_models.models.rfdetr.pre_processing import pre_process_network_input - -try: - from inference_models.models.rfdetr.triton_preprocess import ( - TRITON_AVAILABLE as _TRITON_AVAILABLE, - build_resample_tables, - triton_preprocess_rfdetr_stretch, - ) -except ImportError: - _TRITON_AVAILABLE = False - build_resample_tables = None - triton_preprocess_rfdetr_stretch = None from inference_models.weights_providers.entities import RecommendedParameters try: @@ -99,84 +85,6 @@ ) from import_error -class _FastPathState: - """Per-(src_shape, target_shape) cache of GPU buffers + resample tables - that the Triton fast path reuses across frames.""" - - __slots__ = ( - "src_h", - "src_w", - "target_h", - "target_w", - "pinned_host", - "src_gpu", - "out_buffer", - "tables", - ) - - def __init__( - self, - src_h: int, - src_w: int, - target_h: int, - target_w: int, - pinned_host: torch.Tensor, - src_gpu: torch.Tensor, - out_buffer: torch.Tensor, - tables, - ) -> None: - self.src_h = src_h - self.src_w = src_w - self.target_h = target_h - self.target_w = target_w - self.pinned_host = pinned_host - self.src_gpu = src_gpu - self.out_buffer = out_buffer - self.tables = tables - - @classmethod - def build( - cls, - src_h: int, - src_w: int, - target_h: int, - target_w: int, - device: torch.device, - ) -> "_FastPathState": - pinned_host = torch.empty((src_h, src_w, 3), dtype=torch.uint8, pin_memory=True) - src_gpu = torch.empty((src_h, src_w, 3), dtype=torch.uint8, device=device) - out_buffer = torch.empty( - (1, 3, target_h, target_w), dtype=torch.float32, device=device - ) - tables = build_resample_tables( - src_h=src_h, - src_w=src_w, - target_h=target_h, - target_w=target_w, - device=device, - ) - return cls( - src_h=src_h, - src_w=src_w, - target_h=target_h, - target_w=target_w, - pinned_host=pinned_host, - src_gpu=src_gpu, - out_buffer=out_buffer, - tables=tables, - ) - - def is_stale( - self, src_h: int, src_w: int, target_h: int, target_w: int - ) -> bool: - return ( - self.src_h != src_h - or self.src_w != src_w - or self.target_h != target_h - or self.target_w != target_w - ) - - class RFDetrForInstanceSegmentationTRT( InstanceSegmentationModel[ torch.Tensor, @@ -312,7 +220,6 @@ def __init__( self._inference_stream = torch.cuda.Stream(device=self._device) self._thread_local_storage = threading.local() self.recommended_parameters = recommended_parameters - self._fast_path_state: Optional[_FastPathState] = None @property def class_names(self) -> List[str]: @@ -330,14 +237,6 @@ def pre_process( pre_processing_overrides: Optional[PreProcessingOverrides] = None, **kwargs, ) -> Tuple[torch.Tensor, List[PreProcessingMetadata]]: - fast = self._try_fast_preprocess( - images=images, - input_color_format=input_color_format, - image_size=image_size, - pre_processing_overrides=pre_processing_overrides, - ) - if fast is not None: - return fast with torch.cuda.stream(self._pre_process_stream): pre_processed_images, pre_processing_meta = pre_process_network_input( images=images, @@ -351,128 +250,6 @@ def pre_process( self._pre_process_stream.synchronize() return pre_processed_images, pre_processing_meta - def _try_fast_preprocess( - self, - images, - input_color_format, - image_size, - pre_processing_overrides, - ) -> Optional[Tuple[torch.Tensor, List[PreProcessingMetadata]]]: - if not _TRITON_AVAILABLE: - return None - if image_size is not None: - return None - # pre_processing_overrides can only *disable* transforms; it has no - # "enable" knob. The fast path never applies static_crop / grayscale / - # contrast regardless, so the override flags are irrelevant — we just - # gate on whether the image_pre_processing config itself asks for them. - ipp = self._inference_config.image_pre_processing - if ( - (ipp.static_crop is not None and ipp.static_crop.enabled) - or (ipp.contrast is not None and ipp.contrast.enabled) - or (ipp.grayscale is not None and ipp.grayscale.enabled) - ): - return None - - ni = self._inference_config.network_input - if ni.dataset_version_resize_dimensions is not None: - return None - if ni.input_channels != 3: - return None - if ni.scaling_factor not in (None, 255): - return None - if ni.normalization is None: - return None - # When dataset_version_resize_dimensions is None, the prod path collapses - # non-stretch resize modes to a single PIL stretch as well - # (pre_processing.py:_needs_two_step_resize), so we accept all modes here. - if ni.resize_mode not in ( - ResizeMode.STRETCH_TO, - ResizeMode.LETTERBOX, - ResizeMode.CENTER_CROP, - ResizeMode.LETTERBOX_REFLECT_EDGES, - ): - return None - - if isinstance(images, list): - if len(images) != 1: - return None - candidate = images[0] - else: - candidate = images - if not isinstance(candidate, np.ndarray): - return None - if ( - candidate.dtype != np.uint8 - or candidate.ndim != 3 - or candidate.shape[2] != 3 - ): - return None - - caller_mode = ( - ColorMode(input_color_format) - if input_color_format is not None - else ColorMode.BGR - ) - swap_rb = caller_mode != ni.color_mode - - means, stds = ni.normalization - means_t = (float(means[0]), float(means[1]), float(means[2])) - stds_t = (float(stds[0]), float(stds[1]), float(stds[2])) - target_h = ni.training_input_size.height - target_w = ni.training_input_size.width - orig_h, orig_w = int(candidate.shape[0]), int(candidate.shape[1]) - - state = self._fast_path_state - if state is None or state.is_stale( - src_h=orig_h, - src_w=orig_w, - target_h=target_h, - target_w=target_w, - ): - state = _FastPathState.build( - src_h=orig_h, - src_w=orig_w, - target_h=target_h, - target_w=target_w, - device=self._device, - ) - self._fast_path_state = state - - pinned_np = state.pinned_host.numpy() - np.copyto(pinned_np, candidate, casting="no") - - with torch.cuda.stream(self._pre_process_stream): - state.src_gpu.copy_(state.pinned_host, non_blocking=True) - triton_preprocess_rfdetr_stretch( - src=state.src_gpu, - tables=state.tables, - target_h=target_h, - target_w=target_w, - means=means_t, - stds=stds_t, - swap_rb=swap_rb, - out=state.out_buffer, - ) - state.out_buffer.record_stream(self._pre_process_stream) - self._pre_process_stream.synchronize() - - meta = PreProcessingMetadata( - pad_left=0, - pad_top=0, - pad_right=0, - pad_bottom=0, - original_size=ImageDimensions(width=orig_w, height=orig_h), - size_after_pre_processing=ImageDimensions(width=orig_w, height=orig_h), - inference_size=ImageDimensions(width=target_w, height=target_h), - scale_width=target_w / orig_w, - scale_height=target_h / orig_h, - static_crop_offset=StaticCropOffset( - offset_x=0, offset_y=0, crop_width=orig_w, crop_height=orig_h - ), - ) - return state.out_buffer, [meta] - def forward( self, pre_processed_images: torch.Tensor, From 4e08a512d8790b7e8c8c740890f4cf94f4cf9fd6 Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Sat, 9 May 2026 03:13:17 +0000 Subject: [PATCH 03/24] perf(rfdetr): add static_crop support to Triton fast path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Apply static_crop as a load-time offset in the kernel (+ crop-dims-based resample tables), matching apply_static_crop_to_numpy_image's pixel- coordinate percentage math. Extends fast-path coverage from 35 → 55 of 144 rfdetr integration tests, pass rate unchanged (144/144). --- .../models/rfdetr/pre_processing.py | 64 +++++++++++++++---- .../models/rfdetr/triton_preprocess.py | 31 +++++++-- 2 files changed, 78 insertions(+), 17 deletions(-) diff --git a/inference_models/inference_models/models/rfdetr/pre_processing.py b/inference_models/inference_models/models/rfdetr/pre_processing.py index 8a1120c771..348e7a3abe 100644 --- a/inference_models/inference_models/models/rfdetr/pre_processing.py +++ b/inference_models/inference_models/models/rfdetr/pre_processing.py @@ -135,10 +135,12 @@ def pre_process_network_input( ): return _fast_path_preprocess( image_list=image_list, + image_pre_processing=image_pre_processing, network_input=network_input, target_size=target_size, target_device=target_device, input_color_mode=input_color_mode, + pre_processing_overrides=pre_processing_overrides, ) tensors: List[torch.Tensor] = [] @@ -190,15 +192,25 @@ def _fast_path_eligible( # Kernel runs on CUDA. if target_device.type != "cuda": return False - # Image-level transforms (static_crop / grayscale / contrast) aren't in - # the kernel — predicate rejects if any is enabled in config. + # grayscale / contrast aren't implemented in the kernel; static_crop is + # (as a load-time offset + effective-dims substitution). ipp = image_pre_processing - if ( - (ipp.static_crop is not None and ipp.static_crop.enabled) - or (ipp.contrast is not None and ipp.contrast.enabled) - or (ipp.grayscale is not None and ipp.grayscale.enabled) + if (ipp.contrast is not None and ipp.contrast.enabled) or ( + ipp.grayscale is not None and ipp.grayscale.enabled ): return False + # Honour overrides that force-disable static_crop (True means "disable"). + static_crop_overridden = ( + pre_processing_overrides is not None + and pre_processing_overrides.disable_static_crop is True + ) + if ( + ipp.static_crop is not None + and ipp.static_crop.enabled + and not static_crop_overridden + ): + # We'll apply it in the kernel — don't reject here. + pass # Two-stage dataset-version resize isn't in the kernel. ni = network_input if ni.dataset_version_resize_dimensions is not None: @@ -255,10 +267,12 @@ def _as_hwc_uint8_cuda( def _fast_path_preprocess( image_list: List[Union[np.ndarray, torch.Tensor]], + image_pre_processing: ImagePreProcessing, network_input: NetworkInputDefinition, target_size: ImageDimensions, target_device: torch.device, input_color_mode: Optional[ColorMode], + pre_processing_overrides: Optional[PreProcessingOverrides], ) -> Tuple[torch.Tensor, List[PreProcessingMetadata]]: """Per-image Triton launch + stack. Assumes `_fast_path_eligible` passed.""" means, stds = network_input.normalization @@ -272,12 +286,36 @@ def _fast_path_preprocess( # so the swap runs even on unspecified inputs — we replicate that exactly. swap_rb = input_color_mode != network_input.color_mode + # Resolve whether static_crop is active for this call. Computed once + # because the config + override are call-level. + static_crop_overridden = ( + pre_processing_overrides is not None + and pre_processing_overrides.disable_static_crop is True + ) + crop_cfg = image_pre_processing.static_crop + crop_active = ( + crop_cfg is not None and crop_cfg.enabled and not static_crop_overridden + ) + outs: List[torch.Tensor] = [] metas: List[PreProcessingMetadata] = [] for img in image_list: src_gpu = _as_hwc_uint8_cuda(img, target_device) sh, sw = int(src_gpu.shape[0]), int(src_gpu.shape[1]) - tables = _get_resample_tables(target_device, sh, sw, th, tw) + + if crop_active: + # Matches apply_static_crop_to_numpy_image: percentage-based. + x0 = int(crop_cfg.x_min / 100 * sw) + y0 = int(crop_cfg.y_min / 100 * sh) + x1 = int(crop_cfg.x_max / 100 * sw) + y1 = int(crop_cfg.y_max / 100 * sh) + crop_w = x1 - x0 + crop_h = y1 - y0 + else: + x0 = y0 = 0 + crop_w, crop_h = sw, sh + + tables = _get_resample_tables(target_device, crop_h, crop_w, th, tw) out = triton_preprocess_rfdetr_stretch( src=src_gpu, tables=tables, @@ -286,6 +324,10 @@ def _fast_path_preprocess( means=means_t, stds=stds_t, swap_rb=swap_rb, + crop_offset_y=y0, + crop_offset_x=x0, + crop_h=crop_h, + crop_w=crop_w, ) outs.append(out[0]) # drop leading batch dim to match stack below metas.append( @@ -295,12 +337,12 @@ def _fast_path_preprocess( pad_right=0, pad_bottom=0, original_size=ImageDimensions(width=sw, height=sh), - size_after_pre_processing=ImageDimensions(width=sw, height=sh), + size_after_pre_processing=ImageDimensions(width=crop_w, height=crop_h), inference_size=ImageDimensions(width=tw, height=th), - scale_width=tw / sw, - scale_height=th / sh, + scale_width=tw / crop_w, + scale_height=th / crop_h, static_crop_offset=StaticCropOffset( - offset_x=0, offset_y=0, crop_width=sw, crop_height=sh + offset_x=x0, offset_y=y0, crop_width=crop_w, crop_height=crop_h ), ) ) diff --git a/inference_models/inference_models/models/rfdetr/triton_preprocess.py b/inference_models/inference_models/models/rfdetr/triton_preprocess.py index 0b9d7eaa8e..36b98dc3b4 100644 --- a/inference_models/inference_models/models/rfdetr/triton_preprocess.py +++ b/inference_models/inference_models/models/rfdetr/triton_preprocess.py @@ -108,6 +108,8 @@ def _fused_resize_normalize_kernel( src_w, src_stride_h, src_stride_w, + crop_offset_y, + crop_offset_x, dst_stride_c, dst_stride_h, target_h, @@ -130,7 +132,11 @@ def _fused_resize_normalize_kernel( ): """One kernel per (tile_y, tile_x) over target image. - In : src uint8 HWC (src_h, src_w, 3), source color order. + In : src uint8 HWC (src_h, src_w, 3), source color order. The + resample tables are built against `(crop_h, crop_w)` — the + logical source size after a possible static crop — which the + caller passes as `src_h`/`src_w`. `crop_offset_{y,x}` is the + load-time offset into the raw HWC buffer. Out: dst fp32 CHW (1, 3, target_h, target_w), network color order, (pixel/255 - mean)/std. """ @@ -152,9 +158,9 @@ def _fused_resize_normalize_kernel( vacc_2 = tl.zeros((BLOCK_H, BLOCK_W), dtype=tl.int32) for ky in tl.static_range(KSIZE_Y): - # Source row contributing to each output row in this tile. + # Source row (after static crop) contributing to each output row. sy = ymin + ky - sy_c = tl.maximum(tl.minimum(sy, src_h - 1), 0) + sy_c = tl.maximum(tl.minimum(sy, src_h - 1), 0) + crop_offset_y wy = tl.load(wy_ptr + offs_y * KSIZE_Y + ky, mask=mask_y, other=0) # Horizontal pass for (output_rows_in_tile, output_cols_in_tile): @@ -166,7 +172,7 @@ def _fused_resize_normalize_kernel( for kx in tl.static_range(KSIZE_X): sx = xmin + kx - sx_c = tl.maximum(tl.minimum(sx, src_w - 1), 0) + sx_c = tl.maximum(tl.minimum(sx, src_w - 1), 0) + crop_offset_x wx = tl.load(wx_ptr + offs_x * KSIZE_X + kx, mask=mask_x, other=0) base = sy_c[:, None] * src_stride_h + sx_c[None, :] * src_stride_w p0 = tl.load(src_ptr + base + 0, mask=mask_out, other=0).to(tl.int32) @@ -285,17 +291,26 @@ def triton_preprocess_rfdetr_stretch( means: Tuple[float, float, float] = (0.485, 0.456, 0.406), stds: Tuple[float, float, float] = (0.229, 0.224, 0.225), swap_rb: bool = True, + crop_offset_y: int = 0, + crop_offset_x: int = 0, + crop_h: Optional[int] = None, + crop_w: Optional[int] = None, out: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Fused PIL-exact resize + color swap + normalize. Args: src: uint8 CUDA tensor, shape (H, W, 3), HWC layout. - tables: precomputed int32 resample tables from `build_resample_tables`. + tables: precomputed int32 resample tables sized against the *cropped* + source `(crop_h, crop_w)` → `(target_h, target_w)`. target_h, target_w: output spatial dims. means, stds: normalization in output channel order (R, G, B for network_input.color_mode == 'rgb'). swap_rb: if True, source channel 0 → output B (BGR input, RGB network). + crop_offset_y/_x: load-time offset into `src` for a static crop. 0 + means no crop. + crop_h/_w: effective source dims after crop. Defaults to src dims + when no crop is configured. out: optional preallocated fp32 (1, 3, H, W) CUDA tensor. Returns: @@ -311,7 +326,9 @@ def triton_preprocess_rfdetr_stretch( raise ValueError(f"expected HWC 3-channel, got shape={tuple(src.shape)}") src = src.contiguous() - src_h, src_w = int(src.shape[0]), int(src.shape[1]) + raw_src_h, raw_src_w = int(src.shape[0]), int(src.shape[1]) + src_h = crop_h if crop_h is not None else raw_src_h + src_w = crop_w if crop_w is not None else raw_src_w src_stride_h = int(src.stride(0)) src_stride_w = int(src.stride(1)) @@ -360,6 +377,8 @@ def triton_preprocess_rfdetr_stretch( src_w, src_stride_h, src_stride_w, + int(crop_offset_y), + int(crop_offset_x), dst_stride_c, dst_stride_h, target_h, From 666f96f11fbfc07deda2bb4ae29c12808dd08b12 Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Sat, 9 May 2026 03:22:45 +0000 Subject: [PATCH 04/24] perf(rfdetr): accept CHW-layout uint8 tensors in fast path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit torchvision.io.read_image returns CHW uint8 — 72 test calls in the integration suite arrive in that layout. Mirror _tensor_to_hwc_uint8's CHW heuristic (first dim in {1,3,4} and last dim not in {1,3,4}) and permute to HWC before the kernel. Integration coverage 55 → 113 tests hitting fast path, zero regressions. --- .../models/rfdetr/pre_processing.py | 32 +++++++++++++++---- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/inference_models/inference_models/models/rfdetr/pre_processing.py b/inference_models/inference_models/models/rfdetr/pre_processing.py index 348e7a3abe..7e00d19a06 100644 --- a/inference_models/inference_models/models/rfdetr/pre_processing.py +++ b/inference_models/inference_models/models/rfdetr/pre_processing.py @@ -238,29 +238,47 @@ def _fast_path_eligible( # the per-item check below only needs to handle 3D. for img in image_list: if isinstance(img, np.ndarray): - if img.dtype != np.uint8 or img.ndim != 3 or img.shape[2] != 3: + if img.dtype != np.uint8 or img.ndim != 3: + return False + if img.shape[2] != 3 and not _looks_like_chw(img.shape): return False elif isinstance(img, torch.Tensor): - # Only uint8 HWC. Float tensors keep the existing tensor branch - # (F.resize bilinear *without* PIL's antialias — a caller- - # accepted divergence we don't silently change). - if img.dtype != torch.uint8: + # Only uint8 3-channel images (HWC or CHW). Float tensors keep + # the existing tensor branch (F.resize bilinear *without* PIL's + # antialias — a caller-accepted divergence we don't silently + # change). + if img.dtype != torch.uint8 or img.ndim != 3: return False - if img.ndim != 3 or img.shape[-1] != 3: + if img.shape[-1] != 3 and not _looks_like_chw(img.shape): return False else: return False return True +def _looks_like_chw(shape) -> bool: + """Matches _tensor_to_hwc_uint8's CHW heuristic: first dim is 1/3/4 and + last dim is not. Catches torchvision.io.read_image's CHW output.""" + return ( + len(shape) == 3 + and shape[0] in (1, 3, 4) + and shape[-1] not in (1, 3, 4) + ) + + def _as_hwc_uint8_cuda( img: Union[np.ndarray, torch.Tensor], device: torch.device ) -> torch.Tensor: - """Return a contiguous (H, W, 3) uint8 CUDA tensor, copying if needed.""" + """Return a contiguous (H, W, 3) uint8 CUDA tensor, copying if needed. + Accepts HWC or CHW 3D inputs (CHW is torchvision.io.read_image's layout).""" if isinstance(img, torch.Tensor): + if _looks_like_chw(img.shape): + img = img.permute(1, 2, 0) if img.device != device: img = img.to(device=device, non_blocking=True) return img.contiguous() + if _looks_like_chw(img.shape): + img = np.transpose(img, (1, 2, 0)) t = torch.from_numpy(np.ascontiguousarray(img)) return t.to(device=device, non_blocking=True) From c48b0efd5551cdc3c9600e148b3cd169e6be142c Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Mon, 11 May 2026 06:47:44 +0000 Subject: [PATCH 05/24] perf(rfdetr): add kill switch for Triton preproc fast path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit INFERENCE_MODELS_RFDETR_TRITON_PREPROC_ENABLED (default true). Setting it to false short-circuits _fast_path_eligible 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. Verified on rfdetr-seg-nano: the kernel fires on every eligible call when env=true (5000/5000 on full coco/val2017) and never fires when env=false (0/5000), with byte-identical predictions in both states. e2e on vehicles_312px.mp4 (538 frames, rfdetr-seg-nano TRT): env=true : 99.3 fps env=false: 76.2 fps --- .../inference_models/models/rfdetr/pre_processing.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/inference_models/inference_models/models/rfdetr/pre_processing.py b/inference_models/inference_models/models/rfdetr/pre_processing.py index 7e00d19a06..100027fb13 100644 --- a/inference_models/inference_models/models/rfdetr/pre_processing.py +++ b/inference_models/inference_models/models/rfdetr/pre_processing.py @@ -45,6 +45,7 @@ make_the_value_divisible, pre_process_numpy_image, ) +from inference_models.utils.environment import get_boolean_from_env try: from inference_models.models.rfdetr.triton_preprocess import ( @@ -57,6 +58,12 @@ build_resample_tables = None triton_preprocess_rfdetr_stretch = None +# Kill switch: set INFERENCE_MODELS_RFDETR_TRITON_PREPROC_ENABLED=false to force +# the PIL reference path for every call, regardless of other predicates. +_FAST_PATH_ENABLED = get_boolean_from_env( + "INFERENCE_MODELS_RFDETR_TRITON_PREPROC_ENABLED", default=True +) + # Resample tables are cheap (a few KB of int32 weights) but rebuilding them # is a per-call Python loop we can skip. Key: (device_str, src_h, src_w, th, tw). _RESAMPLE_TABLES_CACHE: Dict[Tuple[str, int, int, int, int], "object"] = {} @@ -187,6 +194,8 @@ def _fast_path_eligible( ) -> bool: """True when every item in `image_list` can go through the fused Triton kernel. Predicate is intentionally conservative — any miss → PIL path.""" + if not _FAST_PATH_ENABLED: + return False if not _TRITON_AVAILABLE: return False # Kernel runs on CUDA. From 7ef57d09f7264fadbcf0a272a5ae6108d7cc37a2 Mon Sep 17 00:00:00 2001 From: Aseem Saxena Date: Mon, 11 May 2026 06:48:03 +0000 Subject: [PATCH 06/24] test(rfdetr): repro scripts for preproc parity + bench MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All scripts are driven by INFERENCE_MODELS_RFDETR_TRITON_PREPROC_ENABLED (true → Triton fast path, false → PIL reference). Each run prints a Triton kernel invocation count so it is visible from the console which path handled each image. Scripts: - parity_triton_vs_pil.py — kernel-vs-PIL fp32 ULP check (20 imgs, direct kernel call; bypasses the model stack) - detection_parity_full.py — 5000-img end-to-end parity driver. Spawns one subprocess per env value (so the module- level env read is re-done), pickles per-image detections + Triton call count, then compares. - parity_env_var.py — same idea at 100 imgs, a quick sanity run. - coco_map.py — bbox + segm mAP via pycocotools; run twice with env=true/false to confirm mAP matches to 4 decimals. - preproc_microbench.py — isolated pre_process() timing at 312² / 720×1280 / 1080×1920. - _fastpath_trace.py — shared instrumentation helper. Patches _fast_path_eligible + _fast_path_preprocess + triton_preprocess_rfdetr_stretch + the two PIL fallbacks and prints per-surface call counts at exit. Used by `python run_with_trace.py