diff --git a/docs/instance_segmentation_mask_resolution.md b/docs/instance_segmentation_mask_resolution.md index 397f4ea2ef..90e880634f 100644 --- a/docs/instance_segmentation_mask_resolution.md +++ b/docs/instance_segmentation_mask_resolution.md @@ -111,8 +111,20 @@ the new opt-in flag. Values from 0 to 1 interpolate model-grid and image-grid dimensions. A lower factor does not always mean a smaller grid or faster execution. `accurate` selects factor 1.0, `fast` selects 0.0, and only `tradeoff` reads `tradeoff_factor`. -Damian's Triton optimization is separate; factors other than 1.0 currently use -the fallback post-processor. +The RF-DETR Triton RLE post-processor supports native and upsampled mask grids +at every factor. Boxes and masks share the encoded grid in `mask_size`, including +for empty and deferred results. It preserves fractional float32 boxes on grids +that differ from the image and rounds to integers only on the image grid. +It uses the same target-size rounding as the reference path. Antialiased downsampling, unsupported preprocessing +transforms, and other existing compatibility limits use the reference post-processor. +The existing Triton enablement flag still controls this path; no execution plan is +required. Sparse-record capacity limits are unchanged, including the existing +deferred-mode overflow error. + +The first request for a new mask grid may compile Triton kernels and initialize +interpolation tables. Warm each expected image shape and mask-resolution factor +before serving latency-sensitive requests. Warm inference measurements exclude +this first-use time and memory cost. Workflow versions v1–v4 retain their compatibility guards. Workflow v5, reduced-mask downstream handling, and full tensor support are deferred to a diff --git a/inference_models/docs/changelog.md b/inference_models/docs/changelog.md index 63c9d8254a..2d1c04d529 100644 --- a/inference_models/docs/changelog.md +++ b/inference_models/docs/changelog.md @@ -72,10 +72,15 @@ COCO `size`. Manually constructed reduced dense detections must supply `image_size` for image-space conversion because it cannot be inferred from the reduced tensor. -- The RF-DETR fused Triton post-processor supports factor `1.0` only. Other - factors use the fallback post-processor so the requested resolution is honored. - Reduced-resolution Triton support and its benchmarking are deferred to a - separate PR. +- The RF-DETR fused Triton RLE post-processor now honors `masks_resolution_factor` + for native and upsampled mask grids, including multiclass and deferred streaming + results. Boxes and masks share the selected grid, with float32 fractional boxes + when that grid differs from the image and final integer rounding otherwise. + Empty and deferred results preserve the same grid and box dtype. Antialiased + downsampling and previously unsupported transforms use the reference path. + Multiclass results respect RF-DETR's query-count cap even when `max_detections` + is larger. Multiclass and deferred selection use global top-k before remapping, + without a four-class-per-query cutoff, and encode each selected query mask once. - The `inference_models` adapter now rejects unknown `mask_decode_mode` values with `InvalidMaskDecodeArgument`; it previously ignored the mode. The legacy backend already rejected invalid modes. Validation also runs without opt-in. diff --git a/inference_models/inference_models/models/rfdetr/triton_postprocess.py b/inference_models/inference_models/models/rfdetr/triton_postprocess.py index 9a810f8d29..92d418b653 100644 --- a/inference_models/inference_models/models/rfdetr/triton_postprocess.py +++ b/inference_models/inference_models/models/rfdetr/triton_postprocess.py @@ -40,6 +40,10 @@ from inference_models.models.base.instance_segmentation import InstanceDetections from inference_models.models.base.types import InstancesRLEMasks from inference_models.models.common.roboflow.model_packages import PreProcessingMetadata +from inference_models.models.common.roboflow.post_processing import ( + finalize_instance_segmentation_boxes, + resolve_mask_target_size, +) from inference_models.models.rfdetr.class_remapping import ClassesReMapping try: @@ -63,8 +67,6 @@ _SPARSE_MAX_ROI_WIDTH = 512 _SPARSE_BLOCK_COLS = 8 _SPARSE_MAX_TOTAL_RUNS = INFERENCE_MODELS_RFDETR_TRITON_POSTPROC_MAX_RUNS -_SPARSE_MAX_CLASSES_PER_QUERY = 4 -_SPARSE_TOPK_MAX_TOTAL_RUNS = _SPARSE_MAX_TOTAL_RUNS * _SPARSE_MAX_CLASSES_PER_QUERY # RF-DETR Seg 2XLarge emits 192x192 masks with 300 queries and COCO class # logits. The sparse path supports that shape by scanning source-mask support # in fixed tiles instead of one giant Triton vector. @@ -128,14 +130,15 @@ def _get_interpolation_weights( raise ValueError("Expected antialiased bilinear interpolation to use 2 taps") indices = torch.zeros((output_size, 2), dtype=torch.int32, device=device) values = torch.zeros((output_size, 2), dtype=torch.float32, device=device) - for output_index in range(output_size): - source_indices = nonzero[output_index].nonzero(as_tuple=True)[0] - indices[output_index, : source_indices.numel()] = source_indices.to( - dtype=torch.int32 - ) - values[output_index, : source_indices.numel()] = weights[ - output_index, source_indices - ] + # nonzero returns coordinates in row-major order. With at most two taps, + # the second coordinate in a row is exactly its second interpolation tap. + # Pack all rows together to avoid a GPU synchronization and several tiny + # launches per output coordinate on the first request for a new grid. + output_indices, source_indices = nonzero.nonzero(as_tuple=True) + tap_indices = torch.zeros_like(output_indices) + tap_indices[1:] = (output_indices[1:] == output_indices[:-1]).long() + indices[output_indices, tap_indices] = source_indices.to(dtype=torch.int32) + values[output_indices, tap_indices] = weights[output_indices, source_indices] cached_value = (indices, values) with _INTERPOLATION_WEIGHT_CACHE_LOCK: @@ -233,15 +236,36 @@ def post_process_single_instance_segmentation_result_to_rle_masks_triton( are supported. Returns ``None`` when the caller should use the reference PyTorch/RLE implementation instead. + Boxes and masks use the grid selected by ``masks_resolution_factor``; + image metadata retains original-image dimensions. Unsupported transforms + and antialiased downsampling use the reference path. + The fast path first emits one candidate per query. If any query has more than one class above threshold, the first pass asks for a retry and the - second pass emits up to ``_SPARSE_MAX_CLASSES_PER_QUERY`` query-class - candidates per query. + second pass selects the global top-k query-class candidates and shares + sparse masks between classes referring to the same query. When ``defer_postprocess_sync`` is set, the function enqueues the metadata, sparse RLE, and DtoH copies, then returns a placeholder whose finalizer does CPU assembly later. That is used by the streaming pipeline to keep the next frame's GPU work moving while Python handles previous detections. + + Args: + image_bboxes: Normalized query boxes in cxcywh format. + image_scores: Sigmoid probabilities for each query and class. + image_masks: Per-query mask logits on the model grid. + image_meta: Preprocessing geometry for this image. + threshold: Confidence cutoff; tensor cutoffs use the reference path. + classes_re_mapping: Mapping from model classes to public class IDs. + max_detections: Optional limit applied after global query-class selection. + defer_postprocess_sync: Return an asynchronous handoff for CPU assembly. + masks_resolution_factor: Native-to-image mask-grid interpolation factor. + + Returns: + Detections or a deferred placeholder, or None for reference fallback. + + Raises: + ValueError: The resolution factor is not finite or outside [0, 1]. """ unsupported_reason = _unsupported_triton_postprocess_reason( image_bboxes=image_bboxes, @@ -262,14 +286,31 @@ def post_process_single_instance_segmentation_result_to_rle_masks_triton( image_scores = image_scores.contiguous() image_bboxes = image_bboxes.contiguous() - image_masks = image_masks.contiguous() + # Both mask kernels consume explicit strides. Keeping the model's view + # avoids copying every query mask before selecting detections. class_mapping = classes_re_mapping.class_mapping.contiguous() num_queries, num_classes = image_scores.shape - if max_detections is None: - max_detections = num_queries + # The reference selects at most num_queries query/class pairs before + # applying max_detections, including in the multiclass retry path. + max_detections = ( + num_queries if max_detections is None else min(max_detections, num_queries) + ) mask_height, mask_width = image_masks.shape[-2:] - output_height = image_meta.original_size.height - output_width = image_meta.original_size.width + image_height, image_width = image_meta.original_size + image_size = (image_height, image_width) + output_height, output_width = resolve_mask_target_size( + mask_height=mask_height, + mask_width=mask_width, + size_after_pre_processing=image_meta.size_after_pre_processing, + masks_resolution_factor=masks_resolution_factor, + ) + denorm_size = image_meta.nonsquare_intermediate_size or image_meta.inference_size + box_scale_width, box_scale_height = image_meta.scale_width, image_meta.scale_height + reduced_grid = (output_height, output_width) != image_size + if reduced_grid: + box_scale_width *= image_width / output_width + box_scale_height *= image_height / output_height + confidence_threshold = float(threshold) # Precompute resize tables outside the hot kernel. The tables are tiny @@ -288,118 +329,29 @@ def post_process_single_instance_segmentation_result_to_rle_masks_triton( ) if defer_postprocess_sync: - # Deferred pipeline mode separates class metadata from query mask RLE: - # every class candidate can reuse the one sparse mask generated for its - # source query, so GPU work remains bounded by the number of queries. - topk_metadata_rows = num_queries * _SPARSE_MAX_CLASSES_PER_QUERY - query_metadata = torch.empty( - (num_queries, _HEADER_SIZE), - dtype=torch.float32, - device=image_scores.device, - ) - class_metadata = torch.empty( - (topk_metadata_rows, _HEADER_SIZE), - dtype=torch.float32, - device=image_scores.device, - ) - records = torch.empty( - (_SPARSE_MAX_TOTAL_RUNS + 1, 3), - dtype=torch.int32, - device=image_scores.device, - ) - source_bounds = _allocate_source_bounds( - rows=num_queries, - mask_height=mask_height, - mask_width=mask_width, - device=image_scores.device, - ) - _select_topk_query_class_metadata_kernel[(num_queries,)]( - image_scores, - image_bboxes, - class_mapping, - class_metadata, - query_metadata, - records, - confidence_threshold, - num_queries, - num_classes, - class_mapping.shape[0], - output_height, - output_width, - BLOCK_CLASSES=triton.next_power_of_2(num_classes), - METADATA_STRIDE=_HEADER_SIZE, - MAX_CLASSES_PER_QUERY=_SPARSE_MAX_CLASSES_PER_QUERY, - FLAG_WRITE_QUERY_METADATA=True, - FLAG_OVERFLOW_CLASSES=False, - ) - _positive_source_bounds_kernel[ - ( - num_queries, - triton.cdiv( - mask_height * mask_width, - _SPARSE_SOURCE_BOUNDS_BLOCK_PIXELS, - ), - ) - ]( - image_masks, - query_metadata, - source_bounds, - mask_height, - mask_width, - image_masks.stride(0), - image_masks.stride(1), - image_masks.stride(2), - BLOCK_PIXELS=_SPARSE_SOURCE_BOUNDS_BLOCK_PIXELS, - METADATA_STRIDE=_HEADER_SIZE, - ) - _sparse_atomic_rle_from_metadata_kernel[ - (num_queries, triton.cdiv(_SPARSE_MAX_ROI_WIDTH, _SPARSE_BLOCK_COLS)) - ]( - image_masks, - source_bounds, - y_idx, - y_weight, - x_idx, - x_weight, - query_metadata, - records, - num_queries, - mask_height, - mask_width, - output_height, - output_width, - image_masks.stride(0), - image_masks.stride(1), - image_masks.stride(2), - BLOCK_OUT_H=triton.next_power_of_2(output_height), - BLOCK_OUT_W=triton.next_power_of_2(output_width), - BLOCK_ROI_H=_BLOCK_ROI_H, - MAX_ROI_WIDTH=_SPARSE_MAX_ROI_WIDTH, - MAX_TOTAL_RUNS=_SPARSE_MAX_TOTAL_RUNS, - METADATA_STRIDE=_HEADER_SIZE, - BLOCK_COLS=_SPARSE_BLOCK_COLS, - ) - outputs_consumed_event = torch.cuda.Event() - outputs_consumed_event.record(torch.cuda.current_stream(image_scores.device)) - class_metadata_host = _acquire_pinned_host_buffer(class_metadata) - records_host = _acquire_pinned_host_buffer(records) - # Pinned buffers let the DtoH copies follow the postprocess kernel on - # the CUDA stream while Python starts preparing later frames. - class_metadata_host.copy_(class_metadata, non_blocking=True) - records_host.copy_(records, non_blocking=True) - done_event = torch.cuda.Event() - done_event.record(torch.cuda.current_stream(image_scores.device)) - return _deferred_instance_detections_from_sparse_query_records( - class_metadata_host=class_metadata_host, - records_host=records_host, - keepalive_tensors=(query_metadata, class_metadata, source_bounds, records), - done_event=done_event, - outputs_consumed_event=outputs_consumed_event, - max_total_runs=_SPARSE_MAX_TOTAL_RUNS, - height=output_height, - width=output_width, + result = _post_process_ranked_query_masks( + image_bboxes=image_bboxes, + image_scores=image_scores, + image_masks=image_masks, + class_mapping=class_mapping, + confidence_threshold=confidence_threshold, + image_size=image_size, + box_transform=( + denorm_size.width, + denorm_size.height, + box_scale_width, + box_scale_height, + ), + output_height=output_height, + output_width=output_width, + y_idx=y_idx, + y_weight=y_weight, + x_idx=x_idx, + x_weight=x_weight, max_detections=max_detections, + defer_postprocess_sync=defer_postprocess_sync, ) + return result # First pass: keep the common case small by selecting only the best class # for each query and emitting sparse RLE runs for those query masks. @@ -431,9 +383,15 @@ def post_process_single_instance_segmentation_result_to_rle_masks_triton( class_mapping.shape[0], output_height, output_width, + denorm_size.width, + denorm_size.height, + box_scale_width, + box_scale_height, + reduced_grid, BLOCK_CLASSES=triton.next_power_of_2(num_classes), METADATA_STRIDE=_HEADER_SIZE, FLAG_MULTICLASS=True, + enable_fp_fusion=False, ) _positive_source_bounds_kernel[ ( @@ -487,6 +445,7 @@ def post_process_single_instance_segmentation_result_to_rle_masks_triton( max_total_runs=_SPARSE_MAX_TOTAL_RUNS, height=output_height, width=output_width, + image_size=image_size, max_detections=max_detections, ) if result is not None: @@ -498,33 +457,76 @@ def post_process_single_instance_segmentation_result_to_rle_masks_triton( ): return None - # Retry only when the first pass detected multiple passing classes for a - # query. This preserves RF-DETR's flat top-k query-class semantics without - # paying the expanded metadata/RLE cost on the usual one-class-per-query - # path. - topk_metadata_rows = num_queries * _SPARSE_MAX_CLASSES_PER_QUERY - metadata = torch.empty( - (topk_metadata_rows, _HEADER_SIZE), - dtype=torch.float32, - device=image_scores.device, + # Retry with global top-k metadata and one mask per selected source query. + result = _post_process_ranked_query_masks( + image_bboxes=image_bboxes, + image_scores=image_scores, + image_masks=image_masks, + class_mapping=class_mapping, + confidence_threshold=confidence_threshold, + image_size=image_size, + box_transform=( + denorm_size.width, + denorm_size.height, + box_scale_width, + box_scale_height, + ), + output_height=output_height, + output_width=output_width, + y_idx=y_idx, + y_weight=y_weight, + x_idx=x_idx, + x_weight=x_weight, + max_detections=max_detections, + defer_postprocess_sync=defer_postprocess_sync, ) + return result + + +def _post_process_ranked_query_masks( + *, + image_bboxes: torch.Tensor, + image_scores: torch.Tensor, + image_masks: torch.Tensor, + class_mapping: torch.Tensor, + confidence_threshold: float, + image_size: Tuple[int, int], + box_transform: Tuple[int, int, float, float], + output_height: int, + output_width: int, + y_idx: torch.Tensor, + y_weight: torch.Tensor, + x_idx: torch.Tensor, + x_weight: torch.Tensor, + max_detections: int, + defer_postprocess_sync: bool, +) -> Optional[InstanceDetections]: + """Select global top-k classes, but interpolate each selected query only once.""" + num_queries, num_classes = image_scores.shape + mask_height, mask_width = image_masks.shape[-2:] + # Match the reference's global selection before remapping/filtering. A + # per-query class cap drops valid detections at low confidence thresholds. + top_scores, top_indices = torch.topk(image_scores.flatten(), num_queries) + query_metadata = torch.empty( + (num_queries, _HEADER_SIZE), dtype=torch.float32, device=image_scores.device + ) + class_metadata = torch.empty_like(query_metadata) records = torch.empty( - (_SPARSE_TOPK_MAX_TOTAL_RUNS + 1, 3), - dtype=torch.int32, - device=image_scores.device, + (_SPARSE_MAX_TOTAL_RUNS + 1, 3), dtype=torch.int32, device=image_scores.device ) source_bounds = _allocate_source_bounds( - rows=topk_metadata_rows, + rows=num_queries, mask_height=mask_height, mask_width=mask_width, device=image_scores.device, ) - _select_topk_query_class_metadata_kernel[(num_queries,)]( - image_scores, + _selected_query_class_metadata_kernel[(num_queries,)]( + top_scores, + top_indices, image_bboxes, class_mapping, - metadata, - metadata, + class_metadata, + query_metadata, records, confidence_threshold, num_queries, @@ -532,20 +534,23 @@ def post_process_single_instance_segmentation_result_to_rle_masks_triton( class_mapping.shape[0], output_height, output_width, - BLOCK_CLASSES=triton.next_power_of_2(num_classes), + *box_transform, + (output_height, output_width) != image_size, + BLOCK_QUERIES=triton.next_power_of_2(num_queries), METADATA_STRIDE=_HEADER_SIZE, - MAX_CLASSES_PER_QUERY=_SPARSE_MAX_CLASSES_PER_QUERY, - FLAG_WRITE_QUERY_METADATA=False, - FLAG_OVERFLOW_CLASSES=True, + enable_fp_fusion=False, ) _positive_source_bounds_kernel[ ( - topk_metadata_rows, - triton.cdiv(mask_height * mask_width, _SPARSE_SOURCE_BOUNDS_BLOCK_PIXELS), + num_queries, + triton.cdiv( + mask_height * mask_width, + _SPARSE_SOURCE_BOUNDS_BLOCK_PIXELS, + ), ) ]( image_masks, - metadata, + query_metadata, source_bounds, mask_height, mask_width, @@ -556,10 +561,7 @@ def post_process_single_instance_segmentation_result_to_rle_masks_triton( METADATA_STRIDE=_HEADER_SIZE, ) _sparse_atomic_rle_from_metadata_kernel[ - ( - topk_metadata_rows, - triton.cdiv(_SPARSE_MAX_ROI_WIDTH, _SPARSE_BLOCK_COLS), - ) + (num_queries, triton.cdiv(_SPARSE_MAX_ROI_WIDTH, _SPARSE_BLOCK_COLS)) ]( image_masks, source_bounds, @@ -567,9 +569,9 @@ def post_process_single_instance_segmentation_result_to_rle_masks_triton( y_weight, x_idx, x_weight, - metadata, + query_metadata, records, - topk_metadata_rows, + num_queries, mask_height, mask_width, output_height, @@ -581,17 +583,42 @@ def post_process_single_instance_segmentation_result_to_rle_masks_triton( BLOCK_OUT_W=triton.next_power_of_2(output_width), BLOCK_ROI_H=_BLOCK_ROI_H, MAX_ROI_WIDTH=_SPARSE_MAX_ROI_WIDTH, - MAX_TOTAL_RUNS=_SPARSE_TOPK_MAX_TOTAL_RUNS, + MAX_TOTAL_RUNS=_SPARSE_MAX_TOTAL_RUNS, METADATA_STRIDE=_HEADER_SIZE, BLOCK_COLS=_SPARSE_BLOCK_COLS, ) - metadata_host = metadata.cpu().numpy() - return _instance_detections_from_sparse_records( - metadata_host=metadata_host, - records=records, - max_total_runs=_SPARSE_TOPK_MAX_TOTAL_RUNS, + if not defer_postprocess_sync: + result = _instance_detections_from_sparse_query_records( + class_metadata_host=class_metadata.cpu().numpy(), + records_host=records.cpu().numpy(), + max_total_runs=_SPARSE_MAX_TOTAL_RUNS, + height=output_height, + width=output_width, + image_size=image_size, + max_detections=max_detections, + ) + return result + + outputs_consumed_event = torch.cuda.Event() + outputs_consumed_event.record(torch.cuda.current_stream(image_scores.device)) + class_metadata_host = _acquire_pinned_host_buffer(class_metadata) + records_host = _acquire_pinned_host_buffer(records) + # Pinned buffers let the DtoH copies follow the postprocess kernel on + # the CUDA stream while Python starts preparing later frames. + class_metadata_host.copy_(class_metadata, non_blocking=True) + records_host.copy_(records, non_blocking=True) + done_event = torch.cuda.Event() + done_event.record(torch.cuda.current_stream(image_scores.device)) + return _deferred_instance_detections_from_sparse_query_records( + class_metadata_host=class_metadata_host, + records_host=records_host, + keepalive_tensors=(query_metadata, class_metadata, source_bounds, records), + done_event=done_event, + outputs_consumed_event=outputs_consumed_event, + max_total_runs=_SPARSE_MAX_TOTAL_RUNS, height=output_height, width=output_width, + image_size=image_size, max_detections=max_detections, ) @@ -603,6 +630,8 @@ def _instance_detections_from_sparse_records( height: int, width: int, max_detections: Optional[int] = None, + *, + image_size: Optional[Tuple[int, int]] = None, ) -> Optional[InstanceDetections]: """Convert sparse device records into ``InstanceDetections``. @@ -613,16 +642,25 @@ def _instance_detections_from_sparse_records( ``None`` means the sparse device result is incomplete or overflowed and the caller should retry or fall back to the reference implementation. """ + image_size = image_size or (height, width) + if np.any(metadata_host[:, 7] > 0.5): + return None + active_ranks = np.flatnonzero(metadata_host[:, 0] > 0.5) if active_ranks.size == 0: return InstanceDetections( - xyxy=torch.empty((0, 4), dtype=torch.int32), + xyxy=finalize_instance_segmentation_boxes( + torch.empty((0, 4)), mask_size=(height, width), image_size=image_size + ), confidence=torch.empty((0,), dtype=torch.float32), class_id=torch.empty((0,), dtype=torch.int32), mask=InstancesRLEMasks.from_coco_rle_masks( - image_size=(height, width), + image_size=image_size, + mask_size=(height, width), masks=[], ), + image_size=image_size, + mask_frame_size=image_size, ) if np.any(metadata_host[active_ranks, 8] > 0.5): return None @@ -643,7 +681,11 @@ def _instance_detections_from_sparse_records( if max_detections is not None: active_ranks = active_ranks[:max_detections] records_host = records_host[1 : total_runs + 1] if total_runs else None - boxes = torch.from_numpy(metadata_host[active_ranks, 3:7].copy()).round().int() + boxes = finalize_instance_segmentation_boxes( + torch.from_numpy(metadata_host[active_ranks, 3:7].copy()), + mask_size=(height, width), + image_size=image_size, + ) confidence = torch.from_numpy(metadata_host[active_ranks, 2].copy()) class_id = torch.from_numpy(metadata_host[active_ranks, 1].copy()).int() @@ -675,7 +717,8 @@ def _instance_detections_from_sparse_records( rle_masks.append(_rle_from_counts(counts=counts, height=height, width=width)) instances_masks = InstancesRLEMasks.from_coco_rle_masks( - image_size=(height, width), + image_size=image_size, + mask_size=(height, width), masks=rle_masks, ) # The pipeline RLE-to-polygon path consumes uncompressed counts directly, so @@ -686,6 +729,8 @@ def _instance_detections_from_sparse_records( confidence=confidence, class_id=class_id, mask=instances_masks, + image_size=image_size, + mask_frame_size=image_size, ) @@ -696,24 +741,32 @@ def _instance_detections_from_sparse_query_records( height: int, width: int, max_detections: Optional[int] = None, + *, + image_size: Optional[Tuple[int, int]] = None, ) -> Optional[InstanceDetections]: """Assemble detections when class rows share query-level RLE records. - Deferred pipeline mode emits up to four class candidates per query, but the - mask is identical for those class rows. The GPU therefore writes RLE records + Multiple selected classes can refer to one query, but the mask is identical + for those class rows. The GPU therefore writes RLE records once per query and this CPU helper fans that query mask out to the selected class detections. """ + image_size = image_size or (height, width) active_ranks = np.flatnonzero(class_metadata_host[:, 0] > 0.5) if active_ranks.size == 0: return InstanceDetections( - xyxy=torch.empty((0, 4), dtype=torch.int32), + xyxy=finalize_instance_segmentation_boxes( + torch.empty((0, 4)), mask_size=(height, width), image_size=image_size + ), confidence=torch.empty((0,), dtype=torch.float32), class_id=torch.empty((0,), dtype=torch.int32), mask=InstancesRLEMasks.from_coco_rle_masks( - image_size=(height, width), + image_size=image_size, + mask_size=(height, width), masks=[], ), + image_size=image_size, + mask_frame_size=image_size, ) if np.any(class_metadata_host[active_ranks, 8] > 0.5): return None @@ -745,8 +798,10 @@ class detections. else: records_host = None record_queries = None - boxes = ( - torch.from_numpy(class_metadata_host[active_ranks, 3:7].copy()).round().int() + boxes = finalize_instance_segmentation_boxes( + torch.from_numpy(class_metadata_host[active_ranks, 3:7].copy()), + mask_size=(height, width), + image_size=image_size, ) confidence = torch.from_numpy(class_metadata_host[active_ranks, 2].copy()) class_id = torch.from_numpy(class_metadata_host[active_ranks, 1].copy()).int() @@ -779,7 +834,8 @@ class detections. rle_masks.append(_rle_from_counts(counts=counts, height=height, width=width)) instances_masks = InstancesRLEMasks.from_coco_rle_masks( - image_size=(height, width), + image_size=image_size, + mask_size=(height, width), masks=rle_masks, ) # The pipeline RLE-to-polygon path consumes uncompressed counts directly, so @@ -790,6 +846,8 @@ class detections. confidence=confidence, class_id=class_id, mask=instances_masks, + image_size=image_size, + mask_frame_size=image_size, ) @@ -826,6 +884,8 @@ def _deferred_instance_detections_from_sparse_query_records( height: int, width: int, max_detections: Optional[int], + *, + image_size: Optional[Tuple[int, int]] = None, ) -> InstanceDetections: """Return a placeholder detection object with deferred CPU finalization. @@ -835,6 +895,8 @@ def _deferred_instance_detections_from_sparse_query_records( ``InstanceDetections``. """ + image_size = image_size or (height, width) + def finalize() -> InstanceDetections: """Synchronize the DtoH copies and build the real detections.""" try: @@ -848,6 +910,7 @@ def finalize() -> InstanceDetections: max_total_runs=max_total_runs, height=height, width=width, + image_size=image_size, max_detections=max_detections, ) if result is None: @@ -865,13 +928,18 @@ def finalize() -> InstanceDetections: _release_pinned_host_buffer(records_host) detections = InstanceDetections( - xyxy=torch.empty((0, 4), dtype=torch.int32), + xyxy=finalize_instance_segmentation_boxes( + torch.empty((0, 4)), mask_size=(height, width), image_size=image_size + ), confidence=torch.empty((0,), dtype=torch.float32), class_id=torch.empty((0,), dtype=torch.int32), mask=InstancesRLEMasks.from_coco_rle_masks( - image_size=(height, width), + image_size=image_size, + mask_size=(height, width), masks=[], ), + image_size=image_size, + mask_frame_size=image_size, ) attach_deferred_postprocess_handoff( detections=detections, @@ -888,9 +956,13 @@ def _should_retry_sparse_topk_metadata( max_total_runs: int, ) -> bool: """Return whether first-pass sparse metadata needs query-class expansion.""" + if np.any(metadata_host[:, 7] > 0.5): + return True + active_ranks = np.flatnonzero(metadata_host[:, 0] > 0.5) if active_ranks.size == 0 or np.any(metadata_host[active_ranks, 8] > 0.5): return False + records_host = records.cpu().numpy() total_runs = int(records_host[0, 0]) if int(records_host[0, 1]) == 0 or total_runs < 0 or total_runs > max_total_runs: @@ -905,6 +977,7 @@ def _supports_triton_postprocess_path( image_meta: PreProcessingMetadata, threshold: Union[float, torch.Tensor], classes_re_mapping: Optional[ClassesReMapping], + masks_resolution_factor: float = 1.0, ) -> bool: """Return ``True`` when the sparse Triton path can represent this input.""" return ( @@ -915,6 +988,7 @@ def _supports_triton_postprocess_path( image_meta=image_meta, threshold=threshold, classes_re_mapping=classes_re_mapping, + masks_resolution_factor=masks_resolution_factor, ) is None ) @@ -932,10 +1006,6 @@ def _unsupported_triton_postprocess_reason( """Explain why the Triton path should not run, or ``None`` when supported.""" if triton is None: return "triton_unavailable" - if masks_resolution_factor != 1.0: - # the fused kernel interpolates straight to the image; it cannot honour - # a reduced target, so defer rather than silently ignore the request - return "mask_resolution_factor_unsupported" if classes_re_mapping is None: return "class_remapping_required" if isinstance(threshold, torch.Tensor): @@ -979,6 +1049,24 @@ def _unsupported_triton_postprocess_reason( or image_meta.size_after_pre_processing.width != output_width ): return "resize_metadata_unsupported" + target_height, target_width = resolve_mask_target_size( + mask_height=mask_height, + mask_width=mask_width, + size_after_pre_processing=image_meta.size_after_pre_processing, + masks_resolution_factor=masks_resolution_factor, + ) + if ( + target_height > 4096 + or target_width > 4096 + or target_height * target_width + > INFERENCE_MODELS_RFDETR_TRITON_POSTPROC_MAX_PIXELS + or target_height * target_width >= _MAX_EXACT_FLAT_INDEX + ): + return "input_size_exceeds_triton_limits" + if target_height < mask_height or target_width < mask_width: + # Antialiased downsampling needs more than the kernel's two taps per + # axis. Keep the reference resize instead of failing in table creation. + return "mask_downsampling_unsupported" if image_scores.device.type != "cuda": return "cuda_device_required" if ( @@ -1084,6 +1172,31 @@ def _attach_uncompressed_counts( if triton is not None: + @triton.jit + def _box_coordinate_on_mask_grid( + center, + extent, + network_size: tl.constexpr, + scale: tl.constexpr, + output_size: tl.constexpr, + lower: tl.constexpr, + reduced_grid: tl.constexpr, + ): + """Mirror reference arithmetic and dtype boundaries before final rounding.""" + half_extent = (0.5 * extent).to(center.dtype) + if lower: + coordinate = (center - half_extent).to(center.dtype) + else: + coordinate = (center + half_extent).to(center.dtype) + coordinate = (coordinate * network_size).to(center.dtype) + if reduced_grid: + divisor = tl.full((), scale, tl.float32) + coordinate = tl.div_rn(coordinate.to(tl.float32), divisor) + else: + divisor = tl.full((), scale, center.dtype).to(tl.float32) + coordinate = tl.div_rn(coordinate.to(tl.float32), divisor).to(center.dtype) + return tl.minimum(tl.maximum(coordinate, 0.0), output_size).to(tl.float32) + @triton.jit def _select_best_query_metadata_kernel( scores, @@ -1097,6 +1210,11 @@ def _select_best_query_metadata_kernel( class_mapping_size: tl.constexpr, output_height: tl.constexpr, output_width: tl.constexpr, + network_width: tl.constexpr, + network_height: tl.constexpr, + box_scale_width: tl.constexpr, + box_scale_height: tl.constexpr, + reduced_grid: tl.constexpr, BLOCK_CLASSES: tl.constexpr, METADATA_STRIDE: tl.constexpr, FLAG_MULTICLASS: tl.constexpr, @@ -1131,15 +1249,18 @@ class is selected. num_classes: Number of model class columns in ``scores``. class_mapping_size: Number of entries available in ``class_mapping``. - output_height: Original image height used to convert normalized box - coordinates to pixel coordinates. - output_width: Original image width used to convert normalized box - coordinates to pixel coordinates. + output_height: Final mask-grid height used to clip boxes. + output_width: Final mask-grid width used to clip boxes. + network_width: Network width for normalized box conversion. + network_height: Network height for normalized box conversion. + box_scale_width: Network-to-mask horizontal divisor. + box_scale_height: Network-to-mask vertical divisor. + reduced_grid: Preserve float32 precision when grids differ. BLOCK_CLASSES: Power-of-two tile width covering ``num_classes``. METADATA_STRIDE: Number of float32 fields per metadata row. - FLAG_MULTICLASS: When true, writes ``records[0, 1] = 1`` if more - than one mapped class for this query exceeds ``threshold`` so - the caller can rerun the top-k query-class path. + FLAG_MULTICLASS: When true, writes ``metadata[rank, 7] = 1`` if more + than one mapped class or any ignored class exceeds threshold, + so the caller can apply global top-k before class filtering. """ rank = tl.program_id(0) meta_base = rank * METADATA_STRIDE @@ -1162,8 +1283,14 @@ class is selected. valid_classes = class_active & (mapped_classes >= 0) passing_classes = valid_classes & (class_scores > threshold) passing_class_count = tl.sum(tl.where(passing_classes, 1, 0), axis=0) - if FLAG_MULTICLASS and passing_class_count > 1: - tl.store(records + 1, 1) + # Ignored classes still consume slots in the reference's global top-k. + # Include inactive queries in this retry decision as well. + ignored_passing = ( + class_active & (mapped_classes < 0) & (class_scores > threshold) + ) + needs_global_selection = (passing_class_count > 1) | ( + tl.sum(ignored_passing.to(tl.int32), axis=0) > 0 + ) # Select over valid mapped classes, not just passing classes. The # threshold is applied after selection so inactive metadata rows still # carry a stable class/score shape. @@ -1196,7 +1323,12 @@ class is selected. metadata + meta_base + 2, tl.where(is_valid_detection, selected_score, 0.0), ) - tl.store(metadata + meta_base + 7, 0.0) + # Each program owns its retry flag; sharing the records header races + # with program zero's initialization on large grids. + tl.store( + metadata + meta_base + 7, + (FLAG_MULTICLASS and needs_global_selection).to(tl.float32), + ) tl.store(metadata + meta_base + 8, 0.0) tl.store(metadata + meta_base + 9, query_index.to(tl.float32)) tl.store(metadata + meta_base + 10, selected_index.to(tl.float32)) @@ -1215,21 +1347,29 @@ class is selected. mask=is_valid_detection, other=0.0, ) - x1 = tl.maximum( - 0.0, - tl.minimum((cx - 0.5 * width) * output_width, output_width), + x1 = _box_coordinate_on_mask_grid( + cx, width, network_width, box_scale_width, output_width, True, reduced_grid ) - y1 = tl.maximum( - 0.0, - tl.minimum((cy - 0.5 * height) * output_height, output_height), + y1 = _box_coordinate_on_mask_grid( + cy, + height, + network_height, + box_scale_height, + output_height, + True, + reduced_grid, ) - x2 = tl.maximum( - 0.0, - tl.minimum((cx + 0.5 * width) * output_width, output_width), + x2 = _box_coordinate_on_mask_grid( + cx, width, network_width, box_scale_width, output_width, False, reduced_grid ) - y2 = tl.maximum( - 0.0, - tl.minimum((cy + 0.5 * height) * output_height, output_height), + y2 = _box_coordinate_on_mask_grid( + cy, + height, + network_height, + box_scale_height, + output_height, + False, + reduced_grid, ) tl.store(metadata + meta_base + 3, x1) tl.store(metadata + meta_base + 4, y1) @@ -1237,8 +1377,9 @@ class is selected. tl.store(metadata + meta_base + 6, y2) @triton.jit - def _select_topk_query_class_metadata_kernel( - scores, + def _selected_query_class_metadata_kernel( + top_scores, + top_indices, bboxes, class_mapping, metadata, @@ -1248,195 +1389,118 @@ def _select_topk_query_class_metadata_kernel( num_queries: tl.constexpr, num_classes: tl.constexpr, class_mapping_size: tl.constexpr, - output_height: tl.constexpr, - output_width: tl.constexpr, - BLOCK_CLASSES: tl.constexpr, + image_height: tl.constexpr, + image_width: tl.constexpr, + network_width: tl.constexpr, + network_height: tl.constexpr, + box_scale_width: tl.constexpr, + box_scale_height: tl.constexpr, + reduced_grid: tl.constexpr, + BLOCK_QUERIES: tl.constexpr, METADATA_STRIDE: tl.constexpr, - MAX_CLASSES_PER_QUERY: tl.constexpr, - FLAG_WRITE_QUERY_METADATA: tl.constexpr, - FLAG_OVERFLOW_CLASSES: tl.constexpr, ): - """Emit top passing query-class metadata rows for one RF-DETR query. + """Write globally selected class rows and a unique active row per query. - Launch grid: - ``(num_queries,)``. Each program scans all class scores for one - query and writes up to ``MAX_CLASSES_PER_QUERY`` rows. The current - implementation uses a static loop of four iterations, so - ``MAX_CLASSES_PER_QUERY`` is expected to be ``4``. - - Args: - scores: CUDA float32 tensor with shape - ``[num_queries, num_classes]`` containing sigmoid class scores. - bboxes: CUDA float32 tensor with shape ``[num_queries, 4]`` in - normalized ``cx, cy, width, height`` format. - class_mapping: CUDA int tensor with class remap entries. Negative - mapped ids are ignored. - metadata: CUDA float32 tensor with shape - ``[num_queries * MAX_CLASSES_PER_QUERY, METADATA_STRIDE]``. - Row ``query_index * MAX_CLASSES_PER_QUERY + class_rank`` holds - the ``class_rank``-th highest passing class for that query. - Columns have the same layout as - ``_select_best_query_metadata_kernel``. - query_metadata: CUDA float32 tensor with shape - ``[num_queries, METADATA_STRIDE]`` when - ``FLAG_WRITE_QUERY_METADATA`` is true. In deferred pipeline - mode, class metadata is expanded but the RLE kernel should still - run once per query, so class rank 0 is also written to this - query-level buffer. When ``FLAG_WRITE_QUERY_METADATA`` is false, - callers pass ``metadata`` here and the argument is unused. - records: CUDA int32 tensor with shape ``[MAX_TOTAL_RUNS + 1, 3]``. - Program 0 resets ``records[0, 0]`` and ``records[0, 1]`` before - the RLE kernel appends runs. - threshold: Minimum class score required for a metadata row to be - marked active. - num_queries: Number of query rows in ``scores`` and ``bboxes``. - num_classes: Number of class columns in ``scores``. - class_mapping_size: Number of valid entries in ``class_mapping``. - output_height: Original image height used for xyxy box conversion. - output_width: Original image width used for xyxy box conversion. - BLOCK_CLASSES: Power-of-two tile width covering all class columns. - METADATA_STRIDE: Number of float32 fields per metadata row. - MAX_CLASSES_PER_QUERY: Number of metadata rows reserved per query. - FLAG_WRITE_QUERY_METADATA: When true, additionally writes the best - passing class row for each query into ``query_metadata`` for the - deferred pipeline RLE kernel. - FLAG_OVERFLOW_CLASSES: When true, writes ``records[0, 1] = 1`` if - more than ``MAX_CLASSES_PER_QUERY`` classes pass threshold; the - caller treats that as unsupported for exact top-k parity. + Each program owns one class row and one query row. Scanning the small + top-k index vector avoids races between classes sharing a source query. + Only selected, valid queries enter the sparse mask kernel. """ - query_index = tl.program_id(0) - if query_index == 0: - tl.store(records + 0, 0) + rank = tl.program_id(0) + if rank == 0: + tl.store(records, 0) tl.store(records + 1, 0) - class_offsets = tl.arange(0, BLOCK_CLASSES) - class_active = class_offsets < num_classes - mapped_classes = tl.load( - class_mapping + class_offsets, - mask=class_active & (class_offsets < class_mapping_size), + score = tl.load(top_scores + rank) + selected_index = tl.load(top_indices + rank).to(tl.int32) + query_index = selected_index // num_classes + class_index = selected_index % num_classes + mapped_class = tl.load( + class_mapping + class_index, + mask=class_index < class_mapping_size, other=-1, ).to(tl.int32) - class_scores = tl.load( - scores + query_index * num_classes + class_offsets, - mask=class_active, - other=-1.0, + active = (score > threshold) & (mapped_class >= 0) + fields = tl.arange(0, METADATA_STRIDE) + tl.store(metadata + rank * METADATA_STRIDE + fields, 0.0) + tl.store(metadata + rank * METADATA_STRIDE, active.to(tl.float32)) + tl.store(metadata + rank * METADATA_STRIDE + 1, mapped_class.to(tl.float32)) + tl.store(metadata + rank * METADATA_STRIDE + 2, score) + tl.store(metadata + rank * METADATA_STRIDE + 9, query_index.to(tl.float32)) + tl.store(metadata + rank * METADATA_STRIDE + 10, selected_index.to(tl.float32)) + cx = tl.load(bboxes + query_index * 4) + cy = tl.load(bboxes + query_index * 4 + 1) + width = tl.load(bboxes + query_index * 4 + 2) + height = tl.load(bboxes + query_index * 4 + 3) + tl.store( + metadata + rank * METADATA_STRIDE + 3, + _box_coordinate_on_mask_grid( + cx, + width, + network_width, + box_scale_width, + image_width, + True, + reduced_grid, + ), ) - passing_classes = ( - class_active & (mapped_classes >= 0) & (class_scores > threshold) + tl.store( + metadata + rank * METADATA_STRIDE + 4, + _box_coordinate_on_mask_grid( + cy, + height, + network_height, + box_scale_height, + image_height, + True, + reduced_grid, + ), ) - passing_class_count = tl.sum(tl.where(passing_classes, 1, 0), axis=0) - if FLAG_OVERFLOW_CLASSES and passing_class_count > MAX_CLASSES_PER_QUERY: - tl.store(records + 1, 1) - - work_scores = tl.where(passing_classes, class_scores, -1.0) - for class_rank in tl.static_range(0, 4): - # Repeated max-and-mask avoids sorting all classes and keeps the - # register footprint bounded by the configured class block. - selected_score = tl.max(work_scores, axis=0) - selected_class = tl.max( - tl.where( - work_scores == selected_score, - class_offsets, - -1, - ), - axis=0, - ).to(tl.int32) - mapped_class = tl.load( - class_mapping + selected_class, - mask=(selected_class >= 0) & (selected_class < class_mapping_size), - other=-1, - ).to(tl.int32) - is_valid_detection = (mapped_class >= 0) & (selected_score > threshold) - metadata_rank = query_index * MAX_CLASSES_PER_QUERY + class_rank - meta_base = metadata_rank * METADATA_STRIDE - selected_index = query_index * num_classes + selected_class - - tl.store( - metadata + meta_base + 0, - tl.where(is_valid_detection, 1.0, 0.0), - ) - tl.store(metadata + meta_base + 1, mapped_class.to(tl.float32)) - tl.store( - metadata + meta_base + 2, - tl.where(is_valid_detection, selected_score, 0.0), - ) - tl.store(metadata + meta_base + 7, 0.0) - tl.store(metadata + meta_base + 8, 0.0) - tl.store(metadata + meta_base + 9, query_index.to(tl.float32)) - tl.store(metadata + meta_base + 10, selected_index.to(tl.float32)) - tl.store(metadata + meta_base + 11, 0.0) - tl.store(metadata + meta_base + 12, 0.0) - tl.store(metadata + meta_base + 13, 0.0) - tl.store(metadata + meta_base + 14, 0.0) - tl.store(metadata + meta_base + 15, 0.0) - - bbox_base = query_index * 4 - cx = tl.load(bboxes + bbox_base, mask=is_valid_detection, other=0.0) - cy = tl.load(bboxes + bbox_base + 1, mask=is_valid_detection, other=0.0) - width = tl.load( - bboxes + bbox_base + 2, - mask=is_valid_detection, - other=0.0, - ) - height = tl.load( - bboxes + bbox_base + 3, - mask=is_valid_detection, - other=0.0, - ) - x1 = tl.maximum( - 0.0, - tl.minimum((cx - 0.5 * width) * output_width, output_width), - ) - y1 = tl.maximum( - 0.0, - tl.minimum((cy - 0.5 * height) * output_height, output_height), - ) - x2 = tl.maximum( - 0.0, - tl.minimum((cx + 0.5 * width) * output_width, output_width), - ) - y2 = tl.maximum( - 0.0, - tl.minimum((cy + 0.5 * height) * output_height, output_height), - ) - tl.store(metadata + meta_base + 3, x1) - tl.store(metadata + meta_base + 4, y1) - tl.store(metadata + meta_base + 5, x2) - tl.store(metadata + meta_base + 6, y2) - if FLAG_WRITE_QUERY_METADATA and class_rank == 0: - # The pipeline path wants best-query metadata for the RLE - # kernel while retaining expanded class metadata for CPU - # finalization. - query_meta_base = query_index * METADATA_STRIDE - tl.store( - query_metadata + query_meta_base + 0, - tl.where(is_valid_detection, 1.0, 0.0), - ) - tl.store( - query_metadata + query_meta_base + 1, mapped_class.to(tl.float32) - ) - tl.store( - query_metadata + query_meta_base + 2, - tl.where(is_valid_detection, selected_score, 0.0), - ) - tl.store(query_metadata + query_meta_base + 3, x1) - tl.store(query_metadata + query_meta_base + 4, y1) - tl.store(query_metadata + query_meta_base + 5, x2) - tl.store(query_metadata + query_meta_base + 6, y2) - tl.store(query_metadata + query_meta_base + 7, 0.0) - tl.store(query_metadata + query_meta_base + 8, 0.0) - tl.store( - query_metadata + query_meta_base + 9, query_index.to(tl.float32) - ) - tl.store( - query_metadata + query_meta_base + 10, selected_index.to(tl.float32) - ) - tl.store(query_metadata + query_meta_base + 11, 0.0) - tl.store(query_metadata + query_meta_base + 12, 0.0) - tl.store(query_metadata + query_meta_base + 13, 0.0) - tl.store(query_metadata + query_meta_base + 14, 0.0) - tl.store(query_metadata + query_meta_base + 15, 0.0) - work_scores = tl.where(class_offsets == selected_class, -1.0, work_scores) + tl.store( + metadata + rank * METADATA_STRIDE + 5, + _box_coordinate_on_mask_grid( + cx, + width, + network_width, + box_scale_width, + image_width, + False, + reduced_grid, + ), + ) + tl.store( + metadata + rank * METADATA_STRIDE + 6, + _box_coordinate_on_mask_grid( + cy, + height, + network_height, + box_scale_height, + image_height, + False, + reduced_grid, + ), + ) + + offsets = tl.arange(0, BLOCK_QUERIES) + indices = tl.load( + top_indices + offsets, mask=offsets < num_queries, other=0 + ).to(tl.int32) + scores = tl.load(top_scores + offsets, mask=offsets < num_queries, other=-1.0) + classes = indices % num_classes + mapped = tl.load( + class_mapping + classes, + mask=(offsets < num_queries) & (classes < class_mapping_size), + other=-1, + ) + selected = ( + (offsets < num_queries) + & (indices // num_classes == rank) + & (scores > threshold) + & (mapped >= 0) + ) + query_active = tl.sum(selected.to(tl.int32), axis=0) > 0 + tl.store(query_metadata + rank * METADATA_STRIDE + fields, 0.0) + tl.store(query_metadata + rank * METADATA_STRIDE, query_active.to(tl.float32)) + tl.store(query_metadata + rank * METADATA_STRIDE + 9, rank.to(tl.float32)) @triton.jit def _positive_source_bounds_kernel( @@ -1576,8 +1640,8 @@ def _sparse_atomic_rle_from_metadata_kernel( num_queries: Number of query masks in ``masks``. mask_height: Height of each low-resolution RF-DETR mask. mask_width: Width of each low-resolution RF-DETR mask. - output_height: Original image height for the output RLE mask. - output_width: Original image width for the output RLE mask. + output_height: Requested mask-grid height for the output RLE mask. + output_width: Requested mask-grid width for the output RLE mask. mask_stride_q: Stride between query masks in ``masks``. mask_stride_h: Row stride for ``masks``. mask_stride_w: Column stride for ``masks``. diff --git a/inference_models/tests/unit_tests/models/rfdetr/test_triton_postprocess.py b/inference_models/tests/unit_tests/models/rfdetr/test_triton_postprocess.py index 83dbc62bbe..0f349d1b8f 100644 --- a/inference_models/tests/unit_tests/models/rfdetr/test_triton_postprocess.py +++ b/inference_models/tests/unit_tests/models/rfdetr/test_triton_postprocess.py @@ -143,6 +143,7 @@ def _expected_result( threshold, classes_re_mapping, num_classes: int = 2, + masks_resolution_factor: float = 1.0, ): return _post_process_single_instance_segmentation_result_to_rle_masks( image_bboxes=bboxes, @@ -152,6 +153,7 @@ def _expected_result( threshold=threshold, num_classes=num_classes, classes_re_mapping=classes_re_mapping, + masks_resolution_factor=masks_resolution_factor, ) @@ -388,10 +390,12 @@ def test_rfdetr_triton_postproc_accepts_2xlarge_shape_limits(monkeypatch) -> Non assert reason == "cuda_device_required" +@pytest.mark.parametrize("factor", [0.0, 0.25, 0.5, 1.0]) @pytest.mark.parametrize("case", ["no_class_mapping", "tensor_threshold", "padding"]) def test_rfdetr_triton_postproc_unsupported_cases_use_reference_path( monkeypatch, case: str, + factor: float, ) -> None: monkeypatch.setattr(rfdetr_common, "_TRITON_POSTPROC_ENABLED", True) calls = 0 @@ -415,7 +419,6 @@ def spy_triton_postprocess(*args, **kwargs): metadata = _metadata() threshold = 0.4 classes_re_mapping = _class_mapping(device) - extra_kwargs = {} if case == "no_class_mapping": classes_re_mapping = None elif case == "tensor_threshold": @@ -430,6 +433,7 @@ def spy_triton_postprocess(*args, **kwargs): metadata=metadata, threshold=threshold, classes_re_mapping=classes_re_mapping, + masks_resolution_factor=factor, ) actual = post_process_instance_segmentation_results_to_rle_masks( bboxes=bboxes.unsqueeze(0), @@ -439,7 +443,7 @@ def spy_triton_postprocess(*args, **kwargs): threshold=threshold, num_classes=2, classes_re_mapping=classes_re_mapping, - **extra_kwargs, + masks_resolution_factor=factor, )[0] assert calls == 1 @@ -673,39 +677,16 @@ def test_rfdetr_triton_postproc_topk_retry_matches_reference_rle_path() -> None: _assert_detections_equal(actual, expected) -@pytest.mark.parametrize( - "factor,expected_supported", [(1.0, True), (0.5, False), (0.0, False)] -) -def test_reduced_mask_resolution_is_unsupported_by_the_fused_path( - monkeypatch, - factor: float, - expected_supported: bool, -) -> None: - # given - # the fused kernel interpolates straight to the image size, so it cannot - # honour a reduced target; it must say so rather than silently ignore it. - # triton is absent on CPU CI, and that reason is checked first, so stand it - # in to reach the check under test +@pytest.mark.parametrize("factor", [0.0, 0.25, 0.5, 1.0]) +def test_reduced_mask_resolution_reaches_the_device_check(monkeypatch, factor) -> None: + # Stand in for Triton on CPU CI to verify that every valid factor passes + # shape/metadata validation before the CUDA requirement is checked. monkeypatch.setattr(triton_postprocess, "triton", object()) - device = torch.device("cpu") - bboxes, logits, masks = _single_detection_inputs(device) - - # when - reason = triton_postprocess._unsupported_triton_postprocess_reason( - image_bboxes=bboxes, - image_scores=logits, - image_masks=masks, - image_meta=_metadata(), - threshold=0.4, - classes_re_mapping=_class_mapping(device), - masks_resolution_factor=factor, + reason = _unsupported_triton_postprocess_reason( + **_support_kwargs(), masks_resolution_factor=factor ) - # then - if expected_supported: - assert reason != "mask_resolution_factor_unsupported" - else: - assert reason == "mask_resolution_factor_unsupported" + assert reason == "cuda_device_required" @pytest.mark.parametrize("factor", [0.25, 0.5]) @@ -742,3 +723,497 @@ def spy(*args, **kwargs): # then assert seen.get("masks_resolution_factor") == factor + + +@pytest.mark.parametrize("shared_queries", [False, True]) +@pytest.mark.parametrize("empty", [False, True]) +@pytest.mark.parametrize("max_detections", [None, 0, 1]) +def test_sparse_assembly_preserves_independent_mask_grid( + shared_queries, empty, max_detections +) -> None: + metadata = np.zeros((2, triton_postprocess._HEADER_SIZE), dtype=np.float32) + if not empty: + metadata[:, 0] = 1 + metadata[:, 1] = [3, 8] + metadata[:, 2] = [0.9, 0.8] + metadata[:, 3:7] = [24, 16, 73, 49] + metadata[:, 9] = 0 # both class rows refer to the same query + records = np.array([[2, 0, 0], [0, 6, 8], [0, 11, 13]], dtype=np.int32) + kwargs = dict( + max_total_runs=2, + height=5, + width=7, + image_size=(65, 97), + max_detections=max_detections, + ) + if shared_queries: + result = triton_postprocess._instance_detections_from_sparse_query_records( + class_metadata_host=metadata, records_host=records, **kwargs + ) + else: + result = triton_postprocess._instance_detections_from_sparse_records( + metadata_host=metadata, records=torch.from_numpy(records), **kwargs + ) + + assert result is not None + assert result.image_size == result.mask_frame_size == (65, 97) + assert result.mask.image_size == (65, 97) + assert result.mask_size == result.mask.mask_size == (5, 7) + count = 0 if empty else min(2, max_detections if max_detections is not None else 2) + assert len(result) == count + if count: + assert result.xyxy.tolist() == [[24, 16, 73, 49]] * count + decoded = coco_rle_masks_to_numpy_mask(result.mask) + expected = np.zeros((5, 7), dtype=bool) + expected[1:3, 1:3] = True + np.testing.assert_array_equal(decoded[0], expected) + if shared_queries and count == 2: + np.testing.assert_array_equal(decoded[1], expected) + + +@pytest.mark.parametrize("empty", [False, True]) +def test_deferred_assembly_preserves_grid_until_finalization( + monkeypatch, empty +) -> None: + from unittest.mock import Mock + + from inference_models.models.base.async_handoff import ( + get_deferred_postprocess_finalizer, + ) + + metadata = torch.zeros((1, triton_postprocess._HEADER_SIZE), dtype=torch.float32) + metadata[0, :7] = torch.tensor([not empty, 3, 0.9, 24, 16, 73, 49]) + records = torch.tensor([[1, 0, 0], [0, 6, 8]], dtype=torch.int32) + done_event = Mock() + release = Mock() + monkeypatch.setattr(triton_postprocess, "_release_pinned_host_buffer", release) + placeholder = ( + triton_postprocess._deferred_instance_detections_from_sparse_query_records( + class_metadata_host=metadata, + records_host=records, + keepalive_tensors=(), + done_event=done_event, + outputs_consumed_event=Mock(), + max_total_runs=1, + height=5, + width=7, + image_size=(65, 97), + max_detections=1, + ) + ) + assert placeholder.mask_size == (5, 7) + assert placeholder.image_size == placeholder.mask_frame_size == (65, 97) + done_event.synchronize.assert_not_called() + result = get_deferred_postprocess_finalizer(placeholder)() + done_event.synchronize.assert_called_once() + assert release.call_count == 2 + assert result.mask_size == (5, 7) + assert result.image_size == result.mask_frame_size == (65, 97) + assert len(result) == (0 if empty else 1) + + +@pytest.mark.parametrize("factor", [-0.1, 1.1, float("nan"), float("inf")]) +def test_triton_rejects_invalid_mask_resolution_factor(monkeypatch, factor) -> None: + monkeypatch.setattr(triton_postprocess, "triton", object()) + with pytest.raises(ValueError, match="finite and in"): + _unsupported_triton_postprocess_reason( + **_support_kwargs(), masks_resolution_factor=factor + ) + + +@pytest.mark.parametrize("factor", [0.0, 0.25, 0.5, 1.0]) +def test_triton_downsampling_uses_the_reference_antialias_path( + monkeypatch, factor +) -> None: + monkeypatch.setattr(triton_postprocess, "triton", object()) + kwargs = _support_kwargs(mask_size=(8, 8)) + kwargs["image_meta"] = _metadata(height=4, width=6) + reason = _unsupported_triton_postprocess_reason( + **kwargs, masks_resolution_factor=factor + ) + assert reason == ( + "cuda_device_required" if factor == 0 else "mask_downsampling_unsupported" + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or triton_postprocess.triton is None, + reason="CUDA and Triton are required", +) +@pytest.mark.parametrize("factor", [0.0, 0.25, 0.5, 1.0]) +@pytest.mark.parametrize("image_size", [(64, 64), (63, 95)]) +@pytest.mark.parametrize( + "mode", ["single", "multiclass", "deferred", "empty", "deferred_empty"] +) +def test_triton_mask_grid_matches_reference_on_cuda(factor, image_size, mode) -> None: + from inference_models.models.base.async_handoff import ( + get_deferred_postprocess_finalizer, + ) + from inference_models.models.common.roboflow.post_processing import ( + resolve_mask_target_size, + ) + + device = torch.device("cuda") + bboxes, logits, masks = _single_detection_inputs(device) + # An asymmetric source and noncontiguous views expose swapped axes/strides. + masks[0, 1:5, 2:7] = 3.0 + masks = masks.transpose(1, 2) + if mode in ("multiclass", "deferred"): + logits[0, 1] = 3.0 + logits[1, 0] = 2.0 + elif mode in ("empty", "deferred_empty"): + logits.fill_(-4.0) + scores = logits.sigmoid().T.contiguous().T + metadata = _metadata(height=image_size[0], width=image_size[1]) + mapping = _class_mapping(device) + mapping.class_mapping[:] = torch.tensor([3, 8], device=device) + expected = _post_process_single_instance_segmentation_result_to_rle_masks( + image_bboxes=bboxes, + image_logits=scores, + image_masks=masks, + image_meta=metadata, + threshold=0.4, + num_classes=2, + classes_re_mapping=mapping, + masks_resolution_factor=factor, + ) + with torch.cuda.stream(torch.cuda.Stream()): + # Producer data were created on the default stream. + torch.cuda.current_stream().wait_stream(torch.cuda.default_stream()) + actual = post_process_single_instance_segmentation_result_to_rle_masks_triton( + image_bboxes=bboxes, + image_scores=scores, + image_masks=masks, + image_meta=metadata, + threshold=0.4, + classes_re_mapping=mapping, + masks_resolution_factor=factor, + defer_postprocess_sync=mode.startswith("deferred"), + max_detections=100, # must still respect RF-DETR's query-count cap + ) + assert actual is not None + if mode.startswith("deferred"): + finalize = get_deferred_postprocess_finalizer(actual) + assert finalize is not None + actual = finalize() + target = resolve_mask_target_size( + 8, + 8, + size_after_pre_processing=metadata.original_size, + masks_resolution_factor=factor, + ) + assert actual.image_size == actual.mask_frame_size == image_size + assert actual.mask.image_size == image_size + assert actual.mask_size == actual.mask.mask_size == target + assert actual.xyxy.dtype == (torch.int32 if target == image_size else torch.float32) + if mode not in ("empty", "deferred_empty"): + _assert_detections_equal(actual, expected) + else: + assert len(actual) == 0 + + +@pytest.mark.skipif( + not torch.cuda.is_available() or triton_postprocess.triton is None, + reason="CUDA and Triton are required", +) +@pytest.mark.parametrize("factor", [0.0, 0.25, 0.5, 1.0]) +def test_triton_native_and_large_source_grids(factor) -> None: + from inference_models.models.common.roboflow.post_processing import ( + resolve_mask_target_size, + ) + + device = torch.device("cuda") + bboxes, logits, _ = _single_detection_inputs(device) + masks = torch.full((2, 192, 192), -2.0, device=device) + masks[0, 12:185, 41:161] = 3.0 + metadata = _metadata(height=1031, width=257) + mapping = _class_mapping(device) + expected = _post_process_single_instance_segmentation_result_to_rle_masks( + image_bboxes=bboxes, + image_logits=logits.sigmoid(), + image_masks=masks, + image_meta=metadata, + threshold=0.4, + num_classes=2, + classes_re_mapping=mapping, + masks_resolution_factor=factor, + ) + actual = post_process_single_instance_segmentation_result_to_rle_masks_triton( + image_bboxes=bboxes, + image_scores=logits.sigmoid(), + image_masks=masks, + image_meta=metadata, + threshold=0.4, + classes_re_mapping=mapping, + masks_resolution_factor=factor, + ) + assert actual is not None + _assert_detections_equal(actual, expected) + assert actual.mask_size == resolve_mask_target_size( + 192, + 192, + size_after_pre_processing=metadata.original_size, + masks_resolution_factor=factor, + ) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or triton_postprocess.triton is None, + reason="CUDA and Triton are required", +) +@pytest.mark.parametrize("factor", [0.0, 0.25, 0.5, 1.0]) +@pytest.mark.parametrize("deferred", [False, True]) +def test_triton_global_topk_keeps_more_than_four_classes_per_query(factor, deferred): + from inference_models.models.base.async_handoff import ( + get_deferred_postprocess_finalizer, + ) + + device = torch.device("cuda") + # The highest eight pairs all come from one query; an ignored class must + # still consume a global top-k slot before remapping is applied. + scores = torch.full((8, 10), 0.01, device=device) + scores[0] = torch.linspace(0.99, 0.80, 10, device=device) + bboxes = torch.tensor([[0.5, 0.5, 0.6, 0.4]], device=device).repeat(8, 1) + masks = torch.full((8, 8, 8), -2.0, device=device) + masks[0, 1:6, 2:7] = 3.0 + mapping = _class_mapping(device, num_classes=10) + mapping.class_mapping[0] = -1 + meta = _metadata(height=63, width=95) + expected = _post_process_single_instance_segmentation_result_to_rle_masks( + image_bboxes=bboxes, + image_logits=scores, + image_masks=masks, + image_meta=meta, + threshold=0.005, + num_classes=10, + classes_re_mapping=mapping, + max_detections=100, + masks_resolution_factor=factor, + ) + actual = post_process_single_instance_segmentation_result_to_rle_masks_triton( + image_bboxes=bboxes, + image_scores=scores, + image_masks=masks, + image_meta=meta, + threshold=0.005, + classes_re_mapping=mapping, + max_detections=100, + masks_resolution_factor=factor, + defer_postprocess_sync=deferred, + ) + assert actual is not None + if deferred: + actual = get_deferred_postprocess_finalizer(actual)() + assert len(actual) == 7 + _assert_detections_equal(actual, expected) + + +@pytest.mark.parametrize("axis", ["height", "width"]) +@pytest.mark.parametrize("source_size, target_size", [(1, 7), (8, 8), (8, 63)]) +def test_batched_interpolation_tables_preserve_reference_weights( + axis, source_size, target_size +): + indices, values = _get_interpolation_weights( + src_size=source_size, + output_size=target_size, + device=torch.device("cpu"), + axis=axis, + ) + basis_shape = ( + (source_size, 1, source_size, 1) + if axis == "height" + else (source_size, 1, 1, source_size) + ) + output_shape = (target_size, 1) if axis == "height" else (1, target_size) + resized = ( + torch.nn.functional.interpolate( + torch.eye(source_size).reshape(basis_shape), + size=output_shape, + mode="bilinear", + align_corners=False, + antialias=True, + ) + .reshape(source_size, target_size) + .T + ) + # Reconstruct every coefficient, including the zero-padded second tap at + # boundaries/native grids, to catch changes to interpolation or tap order. + reconstructed = torch.zeros_like(resized) + reconstructed.scatter_add_(1, indices.long(), values) + torch.testing.assert_close(reconstructed, resized, rtol=0, atol=0) + two_taps = values[:, 1] != 0 + assert torch.all(indices[two_taps, 0] < indices[two_taps, 1]) + assert torch.all(indices[~two_taps, 1] == 0) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or triton_postprocess.triton is None, + reason="CUDA and Triton are required", +) +@pytest.mark.parametrize("factor", [0.0, 0.25, 0.5, 1.0]) +def test_ignored_classes_on_inactive_query_consume_global_topk_slots(factor): + device = torch.device("cuda") + bboxes, _, masks = _single_detection_inputs(device) + scores = torch.tensor([[0.01, 0.02, 0.8], [0.99, 0.9, 0.1]], device=device) + mapping = _class_mapping(device, num_classes=3) + mapping.class_mapping[:] = torch.tensor([-1, -1, 42], device=device) + # Both global top-2 slots are ignored classes on query 1. Query 0's valid + # 0.8 candidate must not be returned, although it passes the confidence cut. + expected = _post_process_single_instance_segmentation_result_to_rle_masks( + image_bboxes=bboxes, + image_logits=scores, + image_masks=masks, + image_meta=_metadata(), + threshold=0.4, + num_classes=3, + classes_re_mapping=mapping, + masks_resolution_factor=factor, + ) + actual = post_process_single_instance_segmentation_result_to_rle_masks_triton( + image_bboxes=bboxes, + image_scores=scores, + image_masks=masks, + image_meta=_metadata(), + threshold=0.4, + classes_re_mapping=mapping, + masks_resolution_factor=factor, + ) + assert actual is not None + assert len(actual) == len(expected) == 0 + # The reference's empty-RLE carrier defaults to the image grid. Triton + # explicitly retains the requested grid even when selection is empty. + expected_side = {0.0: 8, 0.25: 22, 0.5: 36, 1.0: 64}[factor] + assert actual.mask_size == (expected_side, expected_side) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or triton_postprocess.triton is None, + reason="CUDA and Triton are required", +) +@pytest.mark.parametrize("factor", [0.0, 0.25, 0.5, 1.0]) +@pytest.mark.parametrize("deferred", [False, True]) +def test_triton_2xlarge_strided_queries_match_reference(factor, deferred): + from inference_models.models.base.async_handoff import ( + get_deferred_postprocess_finalizer, + ) + + device = torch.device("cuda") + # Exercise the largest supported model grid with a channels-last view, + # 300 global top-k slots, shared masks, and an ignored high-scoring class. + masks = torch.full((192, 192, 300), -2.0, device=device).permute(2, 0, 1) + masks[0, 12:185, 41:161] = 3.0 + masks[1, 40:120, 10:80] = 2.0 + scores = torch.full((300, 91), 0.01, device=device) + scores[0, 1] = 0.95 + scores[0, 2] = 0.90 + scores[1, 3] = 0.85 + scores[2, 0] = 0.99 + bboxes = torch.tensor([[0.5, 0.5, 0.6, 0.4]], device=device).repeat(300, 1) + mapping = _class_mapping(device, num_classes=91) + mapping.class_mapping[0] = -1 + meta = _metadata(height=1031, width=257) + expected = _post_process_single_instance_segmentation_result_to_rle_masks( + image_bboxes=bboxes, + image_logits=scores, + image_masks=masks, + image_meta=meta, + threshold=0.4, + num_classes=91, + classes_re_mapping=mapping, + masks_resolution_factor=factor, + ) + actual = post_process_single_instance_segmentation_result_to_rle_masks_triton( + image_bboxes=bboxes, + image_scores=scores, + image_masks=masks, + image_meta=meta, + threshold=0.4, + classes_re_mapping=mapping, + masks_resolution_factor=factor, + defer_postprocess_sync=deferred, + ) + assert actual is not None + if deferred: + actual = get_deferred_postprocess_finalizer(actual)() + assert len(actual) == 3 + _assert_detections_equal(actual, expected) + + +@pytest.mark.parametrize("shared_queries", [False, True]) +@pytest.mark.parametrize("empty", [False, True]) +@pytest.mark.parametrize("grid", [(5, 7), (65, 97)]) +def test_sparse_boxes_round_only_on_the_image_grid(shared_queries, empty, grid): + metadata = np.zeros((1, triton_postprocess._HEADER_SIZE), dtype=np.float32) + metadata[0, :7] = [not empty, 0, 0.9, 1.25, 2.75, 5.5, 6.25] + records = np.zeros((1, 3), dtype=np.int32) + kwargs = dict(max_total_runs=0, height=grid[0], width=grid[1], image_size=(65, 97)) + if shared_queries: + actual = triton_postprocess._instance_detections_from_sparse_query_records( + class_metadata_host=metadata, records_host=records, **kwargs + ) + else: + actual = triton_postprocess._instance_detections_from_sparse_records( + metadata_host=metadata, records=torch.from_numpy(records), **kwargs + ) + assert actual is not None + expected = torch.tensor([[1.25, 2.75, 5.5, 6.25]])[: 0 if empty else 1] + if grid == (65, 97): + expected = expected.round().int() + torch.testing.assert_close(actual.xyxy, expected, rtol=0, atol=0) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or triton_postprocess.triton is None, + reason="CUDA and Triton are required", +) +@pytest.mark.parametrize("factor", [0.0, 0.5, 1.0]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("mode", ["single", "multiclass", "deferred", "empty"]) +@pytest.mark.parametrize("image_size", [(63, 95), (8, 8)]) +def test_triton_fractional_boxes_match_final_grid_reference( + factor, dtype, mode, image_size +): + from inference_models.models.base.async_handoff import ( + get_deferred_postprocess_finalizer, + ) + + device = torch.device("cuda") + bboxes, logits, masks = _single_detection_inputs(device) + bboxes[0] = torch.tensor([0.501, 0.493, 0.503, 0.489], device=device) + bboxes = bboxes.to(dtype=dtype) + metadata = _metadata(height=image_size[0], width=image_size[1])._replace( + inference_size=ImageDimensions(height=624, width=624), + nonsquare_intermediate_size=ImageDimensions(height=480, width=624), + scale_width=624 / image_size[1], + scale_height=480 / image_size[0], + ) + if mode in ("multiclass", "deferred"): + logits[0, 1] = 3.0 + elif mode == "empty": + logits.fill_(-4.0) + kwargs = dict( + image_bboxes=bboxes, + image_masks=masks, + image_meta=metadata, + threshold=0.4, + classes_re_mapping=_class_mapping(device), + masks_resolution_factor=factor, + ) + scores = logits.sigmoid() + expected = _post_process_single_instance_segmentation_result_to_rle_masks( + image_logits=scores, num_classes=2, **kwargs + ) + actual = post_process_single_instance_segmentation_result_to_rle_masks_triton( + image_scores=scores, defer_postprocess_sync=mode == "deferred", **kwargs + ) + assert actual is not None + if mode == "deferred": + assert actual.xyxy.dtype == expected.xyxy.dtype + actual = get_deferred_postprocess_finalizer(actual)() + assert actual.mask_size == expected.mask_size + torch.testing.assert_close(actual.xyxy.cpu(), expected.xyxy.cpu(), rtol=0, atol=0) + _assert_detections_equal(actual, expected) + restored = actual.to_supervision() + reference_restored = expected.to_supervision() + np.testing.assert_array_equal(restored.xyxy, reference_restored.xyxy) + np.testing.assert_array_equal(restored.mask, reference_restored.mask)