From ab9ab008d213e48c9599355e4fd8913b2f878936 Mon Sep 17 00:00:00 2001 From: yifeif <277870278+yifeif-nv@users.noreply.github.com> Date: Mon, 31 Aug 2026 17:23:03 -0700 Subject: [PATCH 1/6] perf(moge): time TensorRT tactics Move native MoGe builds from optimization level 0 to level 1 and use three timing iterations. This lets TensorRT measure tactics while retaining the full dynamic profile and decomposable FP32 attention needed to avoid the level 3 dynamic-BMM failure. Signed-off-by: yifeif <277870278+yifeif-nv@users.noreply.github.com> --- python/tensorrt_model_connect/families/moge/model.py | 9 +++++---- .../families/moge/tests/test_family.py | 2 ++ 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/python/tensorrt_model_connect/families/moge/model.py b/python/tensorrt_model_connect/families/moge/model.py index 08e49ecfb9..954a57e594 100644 --- a/python/tensorrt_model_connect/families/moge/model.py +++ b/python/tensorrt_model_connect/families/moge/model.py @@ -993,11 +993,12 @@ def _build_native_engine( if hasattr(config, "builder_optimization_level"): # TRT 11.2 optimization level 3 tries to absorb the dynamic FP32 # decomposable-attention chain into one Myelin ForeignNode, whose - # dynamic BMM fallback has no implementation. Level 0 preserves the - # native IAttention decomposition and builds the full 64..2048 profile. - config.builder_optimization_level = 0 + # dynamic BMM fallback has no implementation. Level 1 preserves the + # native IAttention decomposition and full 64..2048 profile while + # still timing tactics instead of accepting the first valid choice. + config.builder_optimization_level = 1 if hasattr(config, "avg_timing_iterations"): - config.avg_timing_iterations = 1 + config.avg_timing_iterations = 3 if hasattr(config, "max_aux_streams"): config.max_aux_streams = 0 if verbose: diff --git a/python/tensorrt_model_connect/families/moge/tests/test_family.py b/python/tensorrt_model_connect/families/moge/tests/test_family.py index ae3207f05b..74b6e3986c 100644 --- a/python/tensorrt_model_connect/families/moge/tests/test_family.py +++ b/python/tensorrt_model_connect/families/moge/tests/test_family.py @@ -122,6 +122,8 @@ def test_production_builder_is_fixed_and_tensor_rt_native_only() -> None: "GELU_ERF", "first_transpose=(0, 3, 1, 2)", "_NUM_TOKENS = 1800", + "config.builder_optimization_level = 1", + "config.avg_timing_iterations = 3", ): assert required in source for output in ("points", "mask", "metric_scale"): From 8f98f0651cade42c515467ee44497866534fadd0 Mon Sep 17 00:00:00 2001 From: yifeif <277870278+yifeif-nv@users.noreply.github.com> Date: Mon, 31 Aug 2026 19:36:21 -0700 Subject: [PATCH 2/6] perf(moge): add 960x540 fast path Add an opt-in FP16 build specialized for 960x540 input. Use native fused attention, rank-2 FP16 dense projections, and an FP16 decoder to remove DRAM-bound attention-score and high-resolution feature traffic while keeping FP32 public outputs and geometry recovery. Preserve the default dynamic FP32 path and fail closed when a fast-path bundle receives another image size. Signed-off-by: yifeif <277870278+yifeif-nv@users.noreply.github.com> --- .../families/moge/model.py | 242 +++++++++++++++--- .../families/moge/tests/test_family.py | 15 +- src/runtime/models/moge/pipeline.cpp | 9 +- src/runtime/models/moge/pipeline.h | 4 +- src/runtime/models/moge/plugin.cpp | 3 +- tests/cpp/models/moge/test_moge_pipeline.cpp | 16 ++ 6 files changed, 239 insertions(+), 50 deletions(-) diff --git a/python/tensorrt_model_connect/families/moge/model.py b/python/tensorrt_model_connect/families/moge/model.py index 954a57e594..92bfefa490 100644 --- a/python/tensorrt_model_connect/families/moge/model.py +++ b/python/tensorrt_model_connect/families/moge/model.py @@ -32,6 +32,8 @@ _MIN_IMAGE_SIZE = 64 _OPT_IMAGE_SIZE = 518 _MAX_IMAGE_SIZE = 2048 +_FAST_IMAGE_HEIGHT = 540 +_FAST_IMAGE_WIDTH = 960 def _require_torch(): @@ -56,10 +58,13 @@ def _checkpoint_digest(path: Path) -> str: class _NativeMogeGraph: """Small family-local vocabulary for composing the exact MoGe graph.""" - def __init__(self, trt: Any, network: Any, state: dict[str, Any]) -> None: + def __init__( + self, trt: Any, network: Any, state: dict[str, Any], *, fast_path: bool = False + ) -> None: self.trt = trt self.network = network self.state = state + self.fast_path = fast_path # TensorRT may retain host weight views until serialization finishes. self._host_weights: list[np.ndarray] = [] @@ -236,11 +241,20 @@ def convolution( *, stride: int = 1, replicate_padding: int = 0, + compute_dtype: Any | None = None, ) -> Any: weight = self._array(f"{module}.weight") if weight.ndim != 4: raise ValueError(f"MoGe convolution {module!r} does not have a 4D kernel") bias = self._array(f"{module}.bias", (weight.shape[0],)) + compute_dtype = compute_dtype or self.trt.float32 + if compute_dtype == self.trt.float16: + weight = np.ascontiguousarray(weight, dtype=np.float16) + bias = np.ascontiguousarray(bias, dtype=np.float16) + self._host_weights.extend((weight, bias)) + elif compute_dtype != self.trt.float32: + raise ValueError(f"Unsupported MoGe convolution compute dtype: {compute_dtype}") + tensor = self.cast(tensor, compute_dtype, f"{name}.input_cast") padded = self.replicate_pad(tensor, replicate_padding, f"{name}.pad") layer = self._layer( self.network.add_convolution_nd( @@ -257,12 +271,27 @@ def convolution( layer.padding_nd = (0, 0) return layer.get_output(0) - def deconvolution(self, tensor: Any, module: str, name: str) -> Any: + def deconvolution( + self, + tensor: Any, + module: str, + name: str, + *, + compute_dtype: Any | None = None, + ) -> Any: weight = self._array(f"{module}.weight") if weight.ndim != 4: raise ValueError(f"MoGe deconvolution {module!r} does not have a 4D kernel") output_channels = int(weight.shape[1]) bias = self._array(f"{module}.bias", (output_channels,)) + compute_dtype = compute_dtype or self.trt.float32 + if compute_dtype == self.trt.float16: + weight = np.ascontiguousarray(weight, dtype=np.float16) + bias = np.ascontiguousarray(bias, dtype=np.float16) + self._host_weights.extend((weight, bias)) + elif compute_dtype != self.trt.float32: + raise ValueError(f"Unsupported MoGe deconvolution compute dtype: {compute_dtype}") + tensor = self.cast(tensor, compute_dtype, f"{name}.input_cast") layer = self._layer( self.network.add_deconvolution_nd( tensor, @@ -278,15 +307,45 @@ def deconvolution(self, tensor: Any, module: str, name: str) -> Any: layer.padding_nd = (0, 0) return layer.get_output(0) - def linear(self, tensor: Any, module: str, name: str) -> Any: + def linear( + self, + tensor: Any, + module: str, + name: str, + *, + compute_dtype: Any | None = None, + output_dtype: Any | None = None, + ) -> Any: weight = self._array(f"{module}.weight") if weight.ndim != 2: raise ValueError(f"MoGe linear {module!r} does not have a 2D weight") output_width, input_width = (int(value) for value in weight.shape) bias = self._array(f"{module}.bias", (output_width,)) + compute_dtype = compute_dtype or self.trt.float32 + if compute_dtype == self.trt.float16: + constant_dtype = np.float16 + weight = np.ascontiguousarray(weight, dtype=np.float16) + bias = np.ascontiguousarray(bias, dtype=np.float16) + self._host_weights.extend((weight, bias)) + elif compute_dtype != self.trt.float32: + raise ValueError(f"Unsupported MoGe linear compute dtype: {compute_dtype}") + else: + constant_dtype = np.float32 + tensor = self.cast(tensor, compute_dtype, f"{name}.input_cast") rank = len(tuple(tensor.shape)) - matrix_shape = (1,) * max(0, rank - 2) + (output_width, input_width) - rhs = self.constant(weight.reshape(matrix_shape), name=f"{name}.weight") + restore_shape = None + if rank > 2: + input_shape = self.shape(tensor, f"{name}.input_shape") + leading = [ + self.shape_index(input_shape, index, f"{name}.output_dim_{index}") + for index in range(rank - 1) + ] + restore_shape = self.shape_concat( + [*leading, self.shape_value(output_width, f"{name}.output_width")], + f"{name}.output_shape", + ) + tensor = self.reshape(tensor, (-1, input_width), f"{name}.input_rows") + rhs = self.constant(weight, dtype=constant_dtype, name=f"{name}.weight") product = self._layer( self.network.add_matrix_multiply( tensor, @@ -297,11 +356,17 @@ def linear(self, tensor: Any, module: str, name: str) -> Any: "matrix multiply", f"{name}.matmul", ).get_output(0) - bias_shape = (1,) * (rank - 1) + (output_width,) - bias_tensor = self.constant(bias.reshape(bias_shape), name=f"{name}.bias") - return self.binary( + bias_tensor = self.constant( + bias.reshape(1, output_width), dtype=constant_dtype, name=f"{name}.bias" + ) + result = self.binary( product, bias_tensor, self.trt.ElementWiseOperation.SUM, f"{name}.bias_add" ) + if restore_shape is not None: + result = self.reshape(result, restore_shape, f"{name}.restore") + if output_dtype is not None: + result = self.cast(result, output_dtype, f"{name}.output_cast") + return result def layer_norm(self, tensor: Any, module: str, name: str) -> Any: rank = len(tuple(tensor.shape)) @@ -438,7 +503,16 @@ def uv(self, height: Any, width: Any, aspect: Any, name: str) -> Any: def attention(self, hidden: Any, layer_index: int, total_tokens: Any) -> Any: prefix = f"encoder.backbone.blocks.{layer_index}.attn" name = f"vit.block.{layer_index}.attention" - qkv = self.linear(hidden, f"{prefix}.qkv", f"{name}.qkv") + if self.fast_path: + qkv = self.linear( + hidden, + f"{prefix}.qkv", + f"{name}.qkv", + compute_dtype=self.trt.float16, + output_dtype=self.trt.float32, + ) + else: + qkv = self.linear(hidden, f"{prefix}.qkv", f"{name}.qkv") component_shape = self.shape_concat( [ self.shape_value(1, f"{name}.batch"), @@ -476,6 +550,10 @@ def attention(self, hidden: Any, layer_index: int, total_tokens: Any) -> Any: ] scale = self.constant([[[[0.125]]]], name=f"{name}.scale") q = self.binary(q, scale, self.trt.ElementWiseOperation.PROD, f"{name}.q_scaled") + if self.fast_path: + q = self.cast(q, self.trt.float16, f"{name}.q_fp16") + k = self.cast(k, self.trt.float16, f"{name}.k_fp16") + v = self.cast(v, self.trt.float16, f"{name}.v_fp16") add_attention_v2 = getattr(self.network, "add_attention_v2", None) if callable(add_attention_v2): layer = add_attention_v2( @@ -490,7 +568,7 @@ def attention(self, hidden: Any, layer_index: int, total_tokens: Any) -> Any: q, k, v, self.trt.AttentionNormalizationOp.SOFTMAX, False ) attention = self._layer(layer, "IAttention", name) - attention.decomposable = True + attention.decomposable = not self.fast_path if hasattr(attention, "query_form"): attention.query_form = self.trt.AttentionIOForm.PADDED_BHND attention.key_value_form = self.trt.AttentionIOForm.PADDED_BHND @@ -508,6 +586,14 @@ def attention(self, hidden: Any, layer_index: int, total_tokens: Any) -> Any: f"{name}.context", first_transpose=(0, 2, 1, 3), ) + if self.fast_path: + return self.linear( + context, + f"{prefix}.proj", + f"{name}.projection", + compute_dtype=self.trt.float16, + output_dtype=self.trt.float32, + ) return self.linear(context, f"{prefix}.proj", f"{name}.projection") def transformer_block(self, hidden: Any, index: int, total_tokens: Any) -> Any: @@ -528,9 +614,26 @@ def transformer_block(self, hidden: Any, index: int, total_tokens: Any) -> Any: hidden, attention, self.trt.ElementWiseOperation.SUM, f"{name}.attention_residual" ) normalized = self.layer_norm(hidden, f"{prefix}.norm2", f"{name}.norm2") - mlp = self.linear(normalized, f"{prefix}.mlp.fc1", f"{name}.mlp.fc1") + if self.fast_path: + mlp = self.linear( + normalized, + f"{prefix}.mlp.fc1", + f"{name}.mlp.fc1", + compute_dtype=self.trt.float16, + ) + else: + mlp = self.linear(normalized, f"{prefix}.mlp.fc1", f"{name}.mlp.fc1") mlp = self.gelu(mlp, f"{name}.mlp.gelu") - mlp = self.linear(mlp, f"{prefix}.mlp.fc2", f"{name}.mlp.fc2") + if self.fast_path: + mlp = self.linear( + mlp, + f"{prefix}.mlp.fc2", + f"{name}.mlp.fc2", + compute_dtype=self.trt.float16, + output_dtype=self.trt.float32, + ) + else: + mlp = self.linear(mlp, f"{prefix}.mlp.fc2", f"{name}.mlp.fc2") gamma2 = self.weight_constant( f"{prefix}.ls2.gamma", expected=(_HIDDEN,), @@ -755,20 +858,37 @@ def encoder(self, image: Any) -> tuple[Any, Any, Any, Any, Any]: class_vector = self.reshape(last_class, (1, _HIDDEN), "vit.class_vector") return encoded, class_vector, base_h, base_w, aspect - def residual_conv_block(self, tensor: Any, module: str, name: str) -> Any: + def residual_conv_block( + self, tensor: Any, module: str, name: str, *, compute_dtype: Any + ) -> Any: hidden = self.relu(tensor, f"{name}.relu1") hidden = self.convolution( - hidden, f"{module}.layers.2", f"{name}.conv1", replicate_padding=1 + hidden, + f"{module}.layers.2", + f"{name}.conv1", + replicate_padding=1, + compute_dtype=compute_dtype, ) hidden = self.relu(hidden, f"{name}.relu2") hidden = self.convolution( - hidden, f"{module}.layers.5", f"{name}.conv2", replicate_padding=1 + hidden, + f"{module}.layers.5", + f"{name}.conv2", + replicate_padding=1, + compute_dtype=compute_dtype, ) return self.binary(hidden, tensor, self.trt.ElementWiseOperation.SUM, f"{name}.residual") - def resample(self, tensor: Any, module: str, level: int, name: str) -> Any: + def resample( + self, tensor: Any, module: str, level: int, name: str, *, compute_dtype: Any + ) -> Any: if level < 3: - tensor = self.deconvolution(tensor, f"{module}.0", f"{name}.deconvolution") + tensor = self.deconvolution( + tensor, + f"{module}.0", + f"{name}.deconvolution", + compute_dtype=compute_dtype, + ) else: shape = self.shape(tensor, f"{name}.input_shape") height = self.shape_index(shape, 2, f"{name}.height") @@ -779,7 +899,13 @@ def resample(self, tensor: Any, module: str, level: int, name: str) -> Any: tensor = self.resize_nchw_to_hw( tensor, output_h, output_w, self.trt.InterpolationMode.LINEAR, f"{name}.resize" ) - return self.convolution(tensor, f"{module}.1", f"{name}.convolution", replicate_padding=1) + return self.convolution( + tensor, + f"{module}.1", + f"{name}.convolution", + replicate_padding=1, + compute_dtype=compute_dtype, + ) def conv_stack( self, @@ -788,12 +914,16 @@ def conv_stack( num_res_blocks: tuple[int, ...], *, final_projection: bool, + compute_dtype: Any, ) -> list[Any]: outputs: list[Any] = [] current = None for level, feature in enumerate(inputs): projected = self.convolution( - feature, f"{prefix}.input_blocks.{level}", f"{prefix}.level.{level}.input" + feature, + f"{prefix}.input_blocks.{level}", + f"{prefix}.level.{level}.input", + compute_dtype=compute_dtype, ) current = ( projected @@ -810,6 +940,7 @@ def conv_stack( current, f"{prefix}.res_blocks.{level}.{block}", f"{prefix}.level.{level}.block.{block}", + compute_dtype=compute_dtype, ) output = current if final_projection and level == len(inputs) - 1: @@ -817,6 +948,7 @@ def conv_stack( current, f"{prefix}.output_blocks.{level}", f"{prefix}.level.{level}.output", + compute_dtype=compute_dtype, ) outputs.append(output) if level < len(inputs) - 1: @@ -825,6 +957,7 @@ def conv_stack( f"{prefix}.resamplers.{level}", level, f"{prefix}.level.{level}.resample", + compute_dtype=compute_dtype, ) return outputs @@ -851,14 +984,27 @@ def outputs(self, image: Any) -> tuple[Any, Any, Any]: else: features.append(uv) + decoder_dtype = self.trt.float16 if self.fast_path else self.trt.float32 neck = self.conv_stack( - features, "neck", (0, 2, 2, 2, 0), final_projection=False + features, + "neck", + (0, 2, 2, 2, 0), + final_projection=False, + compute_dtype=decoder_dtype, ) points = self.conv_stack( - neck, "points_head", (0, 1, 1, 1, 0), final_projection=True + neck, + "points_head", + (0, 1, 1, 1, 0), + final_projection=True, + compute_dtype=decoder_dtype, )[-1] mask = self.conv_stack( - neck, "mask_head", (0, 1, 1, 1, 0), final_projection=True + neck, + "mask_head", + (0, 1, 1, 1, 0), + final_projection=True, + compute_dtype=decoder_dtype, )[-1] scale = self.linear(class_vector, "scale_head.0", "scale_head.0") scale = self.relu(scale, "scale_head.1") @@ -932,8 +1078,10 @@ def outputs(self, image: Any) -> tuple[Any, Any, Any]: def _build_native_engine( state: dict[str, Any], *, + precision: str, verbose: bool, ) -> bytes: + fast_path = precision == "fp16" trt = trt_compat.get_trt() logger = trt.Logger(trt.Logger.INFO if verbose else trt.Logger.WARNING) builder = trt.Builder(logger) @@ -942,10 +1090,11 @@ def _build_native_engine( ) if network is None: raise RuntimeError("TensorRT failed to create the MoGe network") - image = network.add_input("image", trt.float32, (1, -1, -1, 3)) + input_dims = (1, _FAST_IMAGE_HEIGHT, _FAST_IMAGE_WIDTH, 3) if fast_path else (1, -1, -1, 3) + image = network.add_input("image", trt.float32, input_dims) if image is None: raise RuntimeError("TensorRT rejected the MoGe image input") - graph = _NativeMogeGraph(trt, network, state) + graph = _NativeMogeGraph(trt, network, state, fast_path=fast_path) input_shape = graph.shape(image, "input_hwc.shape") input_h = graph.shape_index(input_shape, 1, "input_hwc.height") input_w = graph.shape_index(input_shape, 2, "input_hwc.width") @@ -978,34 +1127,44 @@ def _build_native_engine( config = builder.create_builder_config() tf32 = getattr(trt.BuilderFlag, "TF32", None) if tf32 is not None: - config.clear_flag(tf32) + if fast_path: + config.set_flag(tf32) + else: + config.clear_flag(tf32) config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 16 << 30) profile = builder.create_optimization_profile() - profile.set_shape( - "image", - (1, _MIN_IMAGE_SIZE, _MIN_IMAGE_SIZE, 3), - (1, _OPT_IMAGE_SIZE, _OPT_IMAGE_SIZE, 3), - (1, _MAX_IMAGE_SIZE, _MAX_IMAGE_SIZE, 3), - ) + if fast_path: + fast_shape = (1, _FAST_IMAGE_HEIGHT, _FAST_IMAGE_WIDTH, 3) + profile.set_shape("image", fast_shape, fast_shape, fast_shape) + else: + profile.set_shape( + "image", + (1, _MIN_IMAGE_SIZE, _MIN_IMAGE_SIZE, 3), + (1, _OPT_IMAGE_SIZE, _OPT_IMAGE_SIZE, 3), + (1, _MAX_IMAGE_SIZE, _MAX_IMAGE_SIZE, 3), + ) if not profile: raise RuntimeError("Failed to configure the MoGe dynamic image profile") config.add_optimization_profile(profile) if hasattr(config, "builder_optimization_level"): - # TRT 11.2 optimization level 3 tries to absorb the dynamic FP32 - # decomposable-attention chain into one Myelin ForeignNode, whose - # dynamic BMM fallback has no implementation. Level 1 preserves the - # native IAttention decomposition and full 64..2048 profile while - # still timing tactics instead of accepting the first valid choice. - config.builder_optimization_level = 1 + # Level 3 enables the static FP16 fused-attention fast path. The broad + # dynamic FP32 graph stays at level 1 because level 3 absorbs its + # decomposable attention into an unsupported dynamic-BMM ForeignNode. + config.builder_optimization_level = 3 if fast_path else 1 if hasattr(config, "avg_timing_iterations"): config.avg_timing_iterations = 3 if hasattr(config, "max_aux_streams"): config.max_aux_streams = 0 if verbose: + profile_label = ( + f"{_FAST_IMAGE_WIDTH}x{_FAST_IMAGE_HEIGHT}" + if fast_path + else f"{_MIN_IMAGE_SIZE}..{_MAX_IMAGE_SIZE}" + ) print( "[trtmc build] Building native MoGe TensorRT graph " f"({network.num_layers} layers, num_tokens={_NUM_TOKENS}, " - f"profile={_MIN_IMAGE_SIZE}..{_MAX_IMAGE_SIZE}) ...", + f"precision={precision}, profile={profile_label}) ...", file=sys.stderr, ) plan = builder.build_serialized_network(network, config) @@ -1026,10 +1185,8 @@ def build_moge_engine( checkpoint_path = model_root / _CHECKPOINT if not checkpoint_path.is_file(): raise FileNotFoundError(f"MoGe checkpoint not found: {checkpoint_path}") - if precision != "fp32": - raise ValueError( - "The native MoGe-2 ViT-L accuracy contract supports precision='fp32' only" - ) + if precision not in {"fp32", "fp16"}: + raise ValueError("The native MoGe-2 ViT-L builder supports precision='fp32' or 'fp16' only") checkpoint_sha256 = _checkpoint_digest(checkpoint_path) if checkpoint_sha256 != _CHECKPOINT_SHA256: raise ValueError( @@ -1046,5 +1203,6 @@ def build_moge_engine( ) return _build_native_engine( checkpoint["model"], + precision=precision, verbose=verbose, ) diff --git a/python/tensorrt_model_connect/families/moge/tests/test_family.py b/python/tensorrt_model_connect/families/moge/tests/test_family.py index 74b6e3986c..e043ce0697 100644 --- a/python/tensorrt_model_connect/families/moge/tests/test_family.py +++ b/python/tensorrt_model_connect/families/moge/tests/test_family.py @@ -55,6 +55,7 @@ def test_plugin_keeps_model_state_local_and_rejects_unimplemented_modes( (tmp_path / "model.pt").write_bytes(b"checkpoint") config = SimpleNamespace(raw={}) + assert family_plugin.plugin.default_build_precision == "fp32" assert family_plugin.plugin.load_weights(str(tmp_path), config) == { "model_dir": str(tmp_path.resolve()) } @@ -122,7 +123,11 @@ def test_production_builder_is_fixed_and_tensor_rt_native_only() -> None: "GELU_ERF", "first_transpose=(0, 3, 1, 2)", "_NUM_TOKENS = 1800", - "config.builder_optimization_level = 1", + "_FAST_IMAGE_HEIGHT = 540", + "_FAST_IMAGE_WIDTH = 960", + "attention.decomposable = not self.fast_path", + "compute_dtype=self.trt.float16", + "config.builder_optimization_level = 3 if fast_path else 1", "config.avg_timing_iterations = 3", ): assert required in source @@ -130,10 +135,12 @@ def test_production_builder_is_fixed_and_tensor_rt_native_only() -> None: assert f'("{output}",' in source -def test_build_rejects_unqualified_precision_and_wrong_checkpoint(tmp_path: Path) -> None: +def test_build_rejects_unknown_precision_and_wrong_checkpoint(tmp_path: Path) -> None: (tmp_path / "model.pt").write_bytes(b"wrong checkpoint") - with pytest.raises(ValueError, match="supports precision='fp32' only"): - model_module.build_moge_engine(str(tmp_path), precision="fp16") + with pytest.raises(ValueError, match="supports precision='fp32' or 'fp16' only"): + model_module.build_moge_engine(str(tmp_path), precision="bf16") with pytest.raises(ValueError, match="checkpoint SHA-256 mismatch"): model_module.build_moge_engine(str(tmp_path), precision="fp32") + with pytest.raises(ValueError, match="checkpoint SHA-256 mismatch"): + model_module.build_moge_engine(str(tmp_path), precision="fp16") diff --git a/src/runtime/models/moge/pipeline.cpp b/src/runtime/models/moge/pipeline.cpp index 69f54b2e20..ef754fcf51 100644 --- a/src/runtime/models/moge/pipeline.cpp +++ b/src/runtime/models/moge/pipeline.cpp @@ -355,8 +355,10 @@ moge::GeometryResult postprocess_geometry(const float* affine_points, } // namespace -MogePipeline::MogePipeline(std::unique_ptr model, std::string model_id) - : model_(std::move(model)), model_id_(std::move(model_id)) { +MogePipeline::MogePipeline(std::unique_ptr model, std::string model_id, + bool fixed_fast_shape) + : model_(std::move(model)), model_id_(std::move(model_id)), + fixed_fast_shape_(fixed_fast_shape) { if (!model_ || !model_->ok()) throw std::runtime_error("MogePipeline: invalid model"); } @@ -364,6 +366,9 @@ MogePipeline::MogePipeline(std::unique_ptr model, std::string model_ moge::GeometryResult MogePipeline::estimate_geometry(const float* pixels, int32_t height, int32_t width) { validate_image_input(pixels, height, width); + if (fixed_fast_shape_ && (height != 540 || width != 960)) { + throw std::invalid_argument("MoGe FP16 fast path requires a 960x540 image"); + } Tensor image{const_cast(pixels), {1, height, width, 3}, DType::kFloat32}; const auto outputs = model_->forward({{"image", image}}); diff --git a/src/runtime/models/moge/pipeline.h b/src/runtime/models/moge/pipeline.h index 3046719657..eefbf53848 100644 --- a/src/runtime/models/moge/pipeline.h +++ b/src/runtime/models/moge/pipeline.h @@ -17,7 +17,8 @@ namespace trtmc { class MogePipeline final : public IPipeline, public moge::IGeometryEstimator { public: - MogePipeline(std::unique_ptr model, std::string model_id); + MogePipeline(std::unique_ptr model, std::string model_id, + bool fixed_fast_shape = false); moge::GeometryResult estimate_geometry(const float* pixels, int32_t height, int32_t width) override; @@ -28,6 +29,7 @@ class MogePipeline final : public IPipeline, public moge::IGeometryEstimator { private: std::unique_ptr model_; std::string model_id_; + bool fixed_fast_shape_{false}; }; } // namespace trtmc diff --git a/src/runtime/models/moge/plugin.cpp b/src/runtime/models/moge/plugin.cpp index 15fa9788d5..5292155fc7 100644 --- a/src/runtime/models/moge/plugin.cpp +++ b/src/runtime/models/moge/plugin.cpp @@ -37,7 +37,8 @@ class MogePlugin final : public IPipelinePlugin { public: std::unique_ptr create(const PipelineContext& context) override { return std::make_unique(load_moge_engine(context), - context.bundle.info.model_id); + context.bundle.info.model_id, + context.bundle.info.precision == "fp16"); } }; diff --git a/tests/cpp/models/moge/test_moge_pipeline.cpp b/tests/cpp/models/moge/test_moge_pipeline.cpp index 6eb5ffc88e..e715c1c15a 100644 --- a/tests/cpp/models/moge/test_moge_pipeline.cpp +++ b/tests/cpp/models/moge/test_moge_pipeline.cpp @@ -194,12 +194,28 @@ void test_focal_recovery_failure_is_reported() { } } +void test_fast_path_rejects_non_960x540_input() { + constexpr int32_t size = 64; + auto module = std::make_unique(size, size); + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl", true); + auto image = rgb_image(size, size); + + try { + (void)pipeline.estimate_geometry(image.data(), size, size); + check(false, "MoGe FP16 fast path rejects another shape"); + } catch (const std::invalid_argument& error) { + check(std::string(error.what()).find("960x540") != std::string::npos, + "MoGe FP16 fast-path shape error is explicit"); + } +} + } // namespace int main() { test_pipeline_recovers_metric_geometry_from_hwc_input(); test_invalid_mask_materializes_infinity(); test_focal_recovery_failure_is_reported(); + test_fast_path_rejects_non_960x540_input(); if (g_failures != 0) { std::cerr << g_failures << " MoGe pipeline test(s) failed\n"; return 1; From c2cae4a8673331e15c3c6ea8579aea6bc67e4c68 Mon Sep 17 00:00:00 2001 From: yifeif <277870278+yifeif-nv@users.noreply.github.com> Date: Wed, 2 Sep 2026 02:49:03 -0700 Subject: [PATCH 3/6] perf(moge): accelerate dynamic geometry path Expand the FP16 fast path to the observed dynamic image range while keeping normalization in FP32 and transformer residual and decoder traffic in FP16. Fuse the level-3 resize and convolution, use the qualified family-owned padding selection, and emit compact engine-private geometry outputs. Recover focal length from stable interior neighborhoods while preserving the public geometry contract and unchanged parity thresholds. Keep the implementation native TensorRT and MoGe-owned, without ONNX, FP8, custom plugins, or runtime-core changes. Signed-off-by: yifeif <277870278+yifeif-nv@users.noreply.github.com> --- .../families/moge/model.py | 318 +++++++++++++++--- .../families/moge/tests/test_family.py | 128 ++++++- .../moge/tests/test_fused_resample.py | 240 +++++++++++++ .../moge/tests/test_selective_zero_padding.py | 210 ++++++++++++ src/runtime/models/moge/pipeline.cpp | 110 +++--- src/runtime/models/moge/plugin.cpp | 3 +- tests/cpp/models/moge/test_moge_pipeline.cpp | 181 +++++++++- 7 files changed, 1078 insertions(+), 112 deletions(-) create mode 100644 python/tensorrt_model_connect/families/moge/tests/test_fused_resample.py create mode 100644 python/tensorrt_model_connect/families/moge/tests/test_selective_zero_padding.py diff --git a/python/tensorrt_model_connect/families/moge/model.py b/python/tensorrt_model_connect/families/moge/model.py index 92bfefa490..dc3e50878d 100644 --- a/python/tensorrt_model_connect/families/moge/model.py +++ b/python/tensorrt_model_connect/families/moge/model.py @@ -29,11 +29,66 @@ _PATCH = 14 _POSITION_GRID = 37 _NUM_TOKENS = 1800 +_FOCAL_RECOVERY_SIZE = 64 _MIN_IMAGE_SIZE = 64 _OPT_IMAGE_SIZE = 518 _MAX_IMAGE_SIZE = 2048 -_FAST_IMAGE_HEIGHT = 540 -_FAST_IMAGE_WIDTH = 960 +_FAST_MIN_IMAGE_HEIGHT = 1074 +_FAST_MIN_IMAGE_WIDTH = 608 +_FAST_OPT_IMAGE_HEIGHT = 1080 +_FAST_OPT_IMAGE_WIDTH = 1920 +_FAST_MAX_IMAGE_HEIGHT = 2160 +_FAST_MAX_IMAGE_WIDTH = 3840 +_ZERO_PAD_SELECTION = frozenset( + { + "mask_head.res_blocks.1.0.layers.5", + "mask_head.res_blocks.2.0.layers.2", + "mask_head.res_blocks.2.0.layers.5", + "mask_head.res_blocks.3.0.layers.2", + "mask_head.res_blocks.3.0.layers.5", + "mask_head.resamplers.1.1", + "mask_head.resamplers.2.1", + "neck.res_blocks.2.0.layers.2", + "neck.res_blocks.2.0.layers.5", + "neck.res_blocks.2.1.layers.2", + "neck.res_blocks.2.1.layers.5", + "neck.res_blocks.3.0.layers.2", + "neck.res_blocks.3.0.layers.5", + "neck.res_blocks.3.1.layers.2", + "neck.res_blocks.3.1.layers.5", + "points_head.res_blocks.1.0.layers.2", + "points_head.res_blocks.1.0.layers.5", + "points_head.res_blocks.2.0.layers.2", + "points_head.res_blocks.2.0.layers.5", + "points_head.res_blocks.3.0.layers.2", + "points_head.res_blocks.3.0.layers.5", + "points_head.resamplers.0.1", + "points_head.resamplers.1.1", + "points_head.resamplers.2.1", + } +) + + +def _fuse_half_pixel_x2_conv_weight(weight: np.ndarray) -> np.ndarray: + """Compose HALF_PIXEL bilinear x2 followed by a 3x3 cross-correlation.""" + + if weight.ndim != 4 or tuple(weight.shape[2:]) != (3, 3): + raise ValueError(f"MoGe fused resample requires OI33 weights, got {weight.shape}") + coefficients = np.asarray((0.25, 0.75, 0.75, 0.25), dtype=weight.dtype) + fused = np.zeros((weight.shape[1], weight.shape[0], 6, 6), dtype=weight.dtype) + transposed = weight.transpose(1, 0, 2, 3) + for resize_y, coefficient_y in enumerate(coefficients): + for resize_x, coefficient_x in enumerate(coefficients): + coefficient = coefficient_y * coefficient_x + for kernel_y in range(3): + for kernel_x in range(3): + fused[ + :, + :, + resize_y - kernel_y + 2, + resize_x - kernel_x + 2, + ] += coefficient * transposed[:, :, kernel_y, kernel_x] + return np.ascontiguousarray(fused) def _require_torch(): @@ -185,6 +240,41 @@ def dynamic_slice( layer.mode = mode return layer.get_output(0) + def gather(self, tensor: Any, indices: Any, axis: int, name: str) -> Any: + return self._layer( + self.network.add_gather(tensor, indices, axis), "gather", name + ).get_output(0) + + def nearest_sample_indices(self, size: Any, name: str) -> Any: + positions = self.constant( + np.arange(_FOCAL_RECOVERY_SIZE, dtype=np.int64), + dtype=np.int64, + name=f"{name}.positions", + ) + scaled = self.binary(positions, size, self.trt.ElementWiseOperation.PROD, f"{name}.scaled") + divisor = self.shape_value(_FOCAL_RECOVERY_SIZE, f"{name}.divisor") + return self.binary( + scaled, + divisor, + self.trt.ElementWiseOperation.FLOOR_DIV, + f"{name}.indices", + ) + + def is_finite(self, tensor: Any, name: str) -> Any: + absolute = self.unary(tensor, self.trt.UnaryOperation.ABS, f"{name}.abs") + if tensor.dtype == self.trt.float16: + constant_dtype = np.float16 + elif tensor.dtype == self.trt.float32: + constant_dtype = np.float32 + else: + raise ValueError(f"MoGe finite check {name!r} requires FP16 or FP32 input") + infinity = self.constant( + np.full((1,) * len(tuple(tensor.shape)), np.inf, dtype=constant_dtype), + dtype=constant_dtype, + name=f"{name}.infinity", + ) + return self.binary(absolute, infinity, self.trt.ElementWiseOperation.LESS, name) + def resize( self, tensor: Any, @@ -255,7 +345,19 @@ def convolution( elif compute_dtype != self.trt.float32: raise ValueError(f"Unsupported MoGe convolution compute dtype: {compute_dtype}") tensor = self.cast(tensor, compute_dtype, f"{name}.input_cast") - padded = self.replicate_pad(tensor, replicate_padding, f"{name}.pad") + zero_pad = module in _ZERO_PAD_SELECTION + if zero_pad and ( + replicate_padding != 1 + or stride != 1 + or tuple(int(value) for value in weight.shape[2:]) != (3, 3) + ): + raise ValueError( + f"MoGe selected zero-pad convolution {module!r} must be stride-1 3x3 " + "with one pixel of source replicate padding" + ) + padded = ( + tensor if zero_pad else self.replicate_pad(tensor, replicate_padding, f"{name}.pad") + ) layer = self._layer( self.network.add_convolution_nd( padded, @@ -268,7 +370,7 @@ def convolution( name, ) layer.stride_nd = (stride, stride) - layer.padding_nd = (0, 0) + layer.padding_nd = (1, 1) if zero_pad else (0, 0) return layer.get_output(0) def deconvolution( @@ -307,6 +409,53 @@ def deconvolution( layer.padding_nd = (0, 0) return layer.get_output(0) + def fused_half_pixel_resample( + self, + tensor: Any, + module: str, + name: str, + *, + compute_dtype: Any, + ) -> Any: + """Fuse bilinear x2, replicate padding and a 3x3 convolution exactly.""" + + weight = self._array(f"{module}.weight") + bias = self._array(f"{module}.bias", (int(weight.shape[0]),)) + if compute_dtype == self.trt.float16: + # Match the source convolution's effective checkpoint precision + # before composing its weights with the exact bilinear kernel. + weight = np.ascontiguousarray(weight, dtype=np.float16) + bias = np.ascontiguousarray(bias, dtype=np.float16) + fused_weight = np.ascontiguousarray( + _fuse_half_pixel_x2_conv_weight(weight.astype(np.float32)), + dtype=np.float16, + ) + elif compute_dtype == self.trt.float32: + fused_weight = _fuse_half_pixel_x2_conv_weight(weight) + else: + raise ValueError(f"Unsupported MoGe fused resample dtype: {compute_dtype}") + self._host_weights.extend((weight, bias, fused_weight)) + + tensor = self.cast(tensor, compute_dtype, f"{name}.input_cast") + # Replicating one low-resolution pixel supplies the two high-resolution + # border samples consumed by the original post-resize replicate pad. + tensor = self.replicate_pad(tensor, 1, f"{name}.input_pad") + layer = self._layer( + self.network.add_deconvolution_nd( + tensor, + int(weight.shape[0]), + (6, 6), + self.trt.Weights(fused_weight), + self.trt.Weights(bias), + ), + "fused resample deconvolution", + name, + ) + layer.stride_nd = (2, 2) + # For padded input H+2, (H+2-1)*2 + 6 - 2*4 == 2H. + layer.padding_nd = (4, 4) + return layer.get_output(0) + def linear( self, tensor: Any, @@ -369,6 +518,8 @@ def linear( return result def layer_norm(self, tensor: Any, module: str, name: str) -> Any: + if self.fast_path: + tensor = self.cast(tensor, self.trt.float32, f"{name}.input_fp32") rank = len(tuple(tensor.shape)) width = int(self.state[f"{module}.weight"].numel()) parameter_shape = (1,) * (rank - 1) + (width,) @@ -592,7 +743,7 @@ def attention(self, hidden: Any, layer_index: int, total_tokens: Any) -> Any: f"{prefix}.proj", f"{name}.projection", compute_dtype=self.trt.float16, - output_dtype=self.trt.float32, + output_dtype=self.trt.float16, ) return self.linear(context, f"{prefix}.proj", f"{name}.projection") @@ -607,6 +758,8 @@ def transformer_block(self, hidden: Any, index: int, total_tokens: Any) -> Any: shape=(1, 1, _HIDDEN), name=f"{name}.ls1", ) + if self.fast_path: + gamma1 = self.cast(gamma1, self.trt.float16, f"{name}.ls1_fp16") attention = self.binary( attention, gamma1, self.trt.ElementWiseOperation.PROD, f"{name}.scaled_attention" ) @@ -630,7 +783,7 @@ def transformer_block(self, hidden: Any, index: int, total_tokens: Any) -> Any: f"{prefix}.mlp.fc2", f"{name}.mlp.fc2", compute_dtype=self.trt.float16, - output_dtype=self.trt.float32, + output_dtype=self.trt.float16, ) else: mlp = self.linear(mlp, f"{prefix}.mlp.fc2", f"{name}.mlp.fc2") @@ -640,6 +793,8 @@ def transformer_block(self, hidden: Any, index: int, total_tokens: Any) -> Any: shape=(1, 1, _HIDDEN), name=f"{name}.ls2", ) + if self.fast_path: + gamma2 = self.cast(gamma2, self.trt.float16, f"{name}.ls2_fp16") mlp = self.binary(mlp, gamma2, self.trt.ElementWiseOperation.PROD, f"{name}.scaled_mlp") return self.binary(hidden, mlp, self.trt.ElementWiseOperation.SUM, f"{name}.mlp_residual") @@ -685,6 +840,7 @@ def projected_intermediate( image, f"encoder.output_projections.{projection_index}", f"{name}.projection", + compute_dtype=self.trt.float16 if self.fast_path else self.trt.float32, ) return projected, class_token @@ -758,7 +914,11 @@ def encoder(self, image: Any) -> tuple[Any, Any, Any, Any, Any]: pixels = self.binary(pixels, mean, self.trt.ElementWiseOperation.SUB, "input.center") pixels = self.binary(pixels, std, self.trt.ElementWiseOperation.DIV, "input.normalize") patches = self.convolution( - pixels, "encoder.backbone.patch_embed.proj", "vit.patch_embed", stride=_PATCH + pixels, + "encoder.backbone.patch_embed.proj", + "vit.patch_embed", + stride=_PATCH, + compute_dtype=self.trt.float16 if self.fast_path else self.trt.float32, ) patch_tokens = self.binary( base_h, base_w, self.trt.ElementWiseOperation.PROD, "vit.patch_tokens" @@ -788,6 +948,8 @@ def encoder(self, image: Any) -> tuple[Any, Any, Any, Any, Any]: expected=(1, 1, _HIDDEN), name="vit.class_token", ) + if self.fast_path: + class_token = self.cast(class_token, self.trt.float16, "vit.class_token_fp16") token_concat = self._layer( self.network.add_concatenation([class_token, hidden]), "token concat", "vit.tokens" ) @@ -829,12 +991,17 @@ def encoder(self, image: Any) -> tuple[Any, Any, Any, Any, Any]: "vit.position.tokens", ) position_concat.axis = 1 + position_tokens = position_concat.get_output(0) + if self.fast_path: + position_tokens = self.cast(position_tokens, self.trt.float16, "vit.position_fp16") hidden = self.binary( hidden, - position_concat.get_output(0), + position_tokens, self.trt.ElementWiseOperation.SUM, "vit.tokens_plus_position", ) + if self.fast_path: + hidden = self.cast(hidden, self.trt.float16, "vit.residual_fp16") captured: list[Any] = [] last_class = None @@ -889,6 +1056,13 @@ def resample( f"{name}.deconvolution", compute_dtype=compute_dtype, ) + elif level == 3 and self.fast_path: + return self.fused_half_pixel_resample( + tensor, + f"{module}.1", + f"{name}.fused_deconvolution", + compute_dtype=compute_dtype, + ) else: shape = self.shape(tensor, f"{name}.input_shape") height = self.shape_index(shape, 2, f"{name}.height") @@ -973,6 +1147,8 @@ def outputs(self, image: Any) -> tuple[Any, Any, Any]: base_w, multiplier, self.trt.ElementWiseOperation.PROD, f"uv.level.{level}.width" ) uv = self.uv(height, width, aspect, f"uv.level.{level}") + if self.fast_path: + uv = self.cast(uv, self.trt.float16, f"uv.level.{level}.fp16") if level == 0: concat = self._layer( self.network.add_concatenation([features[0], uv]), @@ -1015,64 +1191,111 @@ def outputs(self, image: Any) -> tuple[Any, Any, Any]: input_shape = self.shape(image, "output.input_shape") input_h = self.shape_index(input_shape, 2, "output.height") input_w = self.shape_index(input_shape, 3, "output.width") - points = self.resize_nchw_to_hw( + raw_points = self.resize_nchw_to_hw( points, input_h, input_w, self.trt.InterpolationMode.LINEAR, "output.points_resize" ) mask = self.resize_nchw_to_hw( mask, input_h, input_w, self.trt.InterpolationMode.LINEAR, "output.mask_resize" ) - points_shape = self.shape_concat( - [ - self.shape_value(1, "output.points_batch"), - input_h, - input_w, - self.shape_value(3, "output.points_channels"), - ], - "output.points_shape", - ) - points = self.reshape( - points, - points_shape, - "output.points_nhwc", - first_transpose=(0, 2, 3, 1), - ) xy_shape = self.shape_concat( [ self.shape_value(1, "output.xy_batch"), + self.shape_value(2, "output.xy_channels"), input_h, input_w, - self.shape_value(2, "output.xy_channels"), ], "output.xy_shape", ) z_shape = self.shape_concat( [ self.shape_value(1, "output.z_batch"), + self.shape_value(1, "output.z_channels"), input_h, input_w, - self.shape_value(1, "output.z_channels"), ], "output.z_shape", ) - xy = self.dynamic_slice(points, (0, 0, 0, 0), xy_shape, "output.xy") - z = self.dynamic_slice(points, (0, 0, 0, 2), z_shape, "output.z") - z = self.unary(z, self.trt.UnaryOperation.EXP, "output.z_exp") - xy = self.binary(xy, z, self.trt.ElementWiseOperation.PROD, "output.xy_scaled") - point_concat = self._layer( - self.network.add_concatenation([xy, z]), "point concat", "output.points_remap" - ) - point_concat.axis = 3 - points = point_concat.get_output(0) + raw_xy = self.dynamic_slice(raw_points, (0, 0, 0, 0), xy_shape, "output.raw_xy") + raw_z = self.dynamic_slice(raw_points, (0, 2, 0, 0), z_shape, "output.raw_z") + z = self.unary(raw_z, self.trt.UnaryOperation.EXP, "output.z_exp") + xy = self.binary(raw_xy, z, self.trt.ElementWiseOperation.PROD, "output.xy_scaled") mask_shape = self.shape_concat( [self.shape_value(1, "output.mask_batch"), input_h, input_w], "output.mask_shape", ) + affine_depth = self.reshape(z, mask_shape, "output.affine_depth_squeeze") + affine_depth = self.cast(affine_depth, self.trt.float32, "output.affine_depth_fp32") + + row_indices = self.nearest_sample_indices(input_h, "output.focal_rows") + column_indices = self.nearest_sample_indices(input_w, "output.focal_columns") + sampled_xy = self.gather(xy, row_indices, 2, "output.focal_xy_rows") + sampled_xy = self.gather(sampled_xy, column_indices, 3, "output.focal_xy_columns") + sampled_z = self.gather(z, row_indices, 2, "output.focal_z_rows") + sampled_z = self.gather(sampled_z, column_indices, 3, "output.focal_z_columns") + sampled_concat = self._layer( + self.network.add_concatenation([sampled_xy, sampled_z]), + "sampled point concat", + "output.focal_samples_nchw", + ) + sampled_concat.axis = 1 + focal_samples = self.reshape( + sampled_concat.get_output(0), + (1, _FOCAL_RECOVERY_SIZE, _FOCAL_RECOVERY_SIZE, 3), + "output.focal_samples_nhwc", + first_transpose=(0, 2, 3, 1), + ) + focal_samples = self.cast(focal_samples, self.trt.float32, "output.focal_samples_fp32") + + x = self.dynamic_slice(xy, (0, 0, 0, 0), z_shape, "output.valid.x") + y = self.dynamic_slice(xy, (0, 1, 0, 0), z_shape, "output.valid.y") + x_finite = self.is_finite(x, "output.valid.x_finite") + y_finite = self.is_finite(y, "output.valid.y_finite") + z_finite = self.is_finite(z, "output.valid.z_finite") + points_finite = self.binary( + x_finite, + y_finite, + self.trt.ElementWiseOperation.AND, + "output.valid.xy_finite", + ) + points_finite = self.binary( + points_finite, + z_finite, + self.trt.ElementWiseOperation.AND, + "output.valid.xyz_finite", + ) + points_finite = self.reshape(points_finite, mask_shape, "output.valid.points_squeeze") + + # Keep the legacy sigmoid and FP16->FP32 boundary. For a tiny positive + # FP16 logit, sigmoid can round to exactly 0.5, so logit > 0 is not an + # exact replacement for the public mask predicate. mask = self.reshape(mask, mask_shape, "output.mask_squeeze") mask = self.sigmoid(mask, "output.mask_sigmoid") + mask = self.cast(mask, self.trt.float32, "output.mask_fp32") + mask_threshold = self.constant( + [[[0.5]]], dtype=np.float32, name="output.valid.mask_threshold" + ) + mask_selected = self.binary( + mask, + mask_threshold, + self.trt.ElementWiseOperation.GREATER, + "output.valid.mask_selected", + ) + # GREATER is ordered: NaN compares false. Sigmoid maps finite values + # and +/-infinity to finite probabilities, so the legacy isfinite(mask) + # term cannot reject anything that this comparison would select. + valid = self.binary( + points_finite, + mask_selected, + self.trt.ElementWiseOperation.AND, + "output.valid.selected", + ) + valid = self.cast(valid, self.trt.float16, "output.valid_fp16") + scale = self.unary(scale, self.trt.UnaryOperation.EXP, "output.metric_scale_exp") scale = self.reshape(scale, (1,), "output.metric_scale_squeeze") - return points, mask, scale + scale = self.cast(scale, self.trt.float32, "output.metric_scale_fp32") + return affine_depth, valid, focal_samples, scale def _build_native_engine( @@ -1090,7 +1313,7 @@ def _build_native_engine( ) if network is None: raise RuntimeError("TensorRT failed to create the MoGe network") - input_dims = (1, _FAST_IMAGE_HEIGHT, _FAST_IMAGE_WIDTH, 3) if fast_path else (1, -1, -1, 3) + input_dims = (1, -1, -1, 3) image = network.add_input("image", trt.float32, input_dims) if image is None: raise RuntimeError("TensorRT rejected the MoGe image input") @@ -1113,14 +1336,13 @@ def _build_native_engine( "input_hwc.to_nchw", first_transpose=(0, 3, 1, 2), ) - points, mask, metric_scale = graph.outputs(image) + affine_depth, valid, focal_samples, metric_scale = graph.outputs(image) for name, tensor in ( - ("points", points), - ("mask", mask), + ("affine_depth", affine_depth), + ("valid", valid), + ("focal_samples", focal_samples), ("metric_scale", metric_scale), ): - if tensor.dtype != trt.float32: - tensor = graph.cast(tensor, trt.float32, f"output.{name}_fp32") tensor.name = name network.mark_output(tensor) @@ -1134,8 +1356,12 @@ def _build_native_engine( config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 16 << 30) profile = builder.create_optimization_profile() if fast_path: - fast_shape = (1, _FAST_IMAGE_HEIGHT, _FAST_IMAGE_WIDTH, 3) - profile.set_shape("image", fast_shape, fast_shape, fast_shape) + profile.set_shape( + "image", + (1, _FAST_MIN_IMAGE_HEIGHT, _FAST_MIN_IMAGE_WIDTH, 3), + (1, _FAST_OPT_IMAGE_HEIGHT, _FAST_OPT_IMAGE_WIDTH, 3), + (1, _FAST_MAX_IMAGE_HEIGHT, _FAST_MAX_IMAGE_WIDTH, 3), + ) else: profile.set_shape( "image", @@ -1157,7 +1383,9 @@ def _build_native_engine( config.max_aux_streams = 0 if verbose: profile_label = ( - f"{_FAST_IMAGE_WIDTH}x{_FAST_IMAGE_HEIGHT}" + f"{_FAST_MIN_IMAGE_WIDTH}x{_FAST_MIN_IMAGE_HEIGHT}.." + f"{_FAST_MAX_IMAGE_WIDTH}x{_FAST_MAX_IMAGE_HEIGHT}" + f"@{_FAST_OPT_IMAGE_WIDTH}x{_FAST_OPT_IMAGE_HEIGHT}" if fast_path else f"{_MIN_IMAGE_SIZE}..{_MAX_IMAGE_SIZE}" ) diff --git a/python/tensorrt_model_connect/families/moge/tests/test_family.py b/python/tensorrt_model_connect/families/moge/tests/test_family.py index e043ce0697..9578b32b03 100644 --- a/python/tensorrt_model_connect/families/moge/tests/test_family.py +++ b/python/tensorrt_model_connect/families/moge/tests/test_family.py @@ -11,6 +11,7 @@ from types import SimpleNamespace import tomllib +import numpy as np import pytest from tensorrt_model_connect import engine_builder @@ -49,7 +50,7 @@ def test_config_adapter_claims_one_flat_checkpoint(tmp_path: Path) -> None: assert not family_plugin.plugin.matches("MoGeModel") -def test_plugin_keeps_model_state_local_and_rejects_unimplemented_modes( +def test_plugin_keeps_model_state_local_and_rejects_unsupported_quantization( tmp_path: Path, ) -> None: (tmp_path / "model.pt").write_bytes(b"checkpoint") @@ -108,6 +109,10 @@ def test_production_builder_is_fixed_and_tensor_rt_native_only() -> None: "add_plugin", "get_plugin_registry", "trtmc_moge_", + "add_quantize", + "add_dequantize", + "fp8_scale_map", + "_fp8_dense_selection", ): assert forbidden not in lowered for required in ( @@ -123,18 +128,133 @@ def test_production_builder_is_fixed_and_tensor_rt_native_only() -> None: "GELU_ERF", "first_transpose=(0, 3, 1, 2)", "_NUM_TOKENS = 1800", - "_FAST_IMAGE_HEIGHT = 540", - "_FAST_IMAGE_WIDTH = 960", + "_FOCAL_RECOVERY_SIZE = 64", + "_FAST_MIN_IMAGE_HEIGHT = 1074", + "_FAST_MIN_IMAGE_WIDTH = 608", + "_FAST_OPT_IMAGE_HEIGHT = 1080", + "_FAST_OPT_IMAGE_WIDTH = 1920", + "_FAST_MAX_IMAGE_HEIGHT = 2160", + "_FAST_MAX_IMAGE_WIDTH = 3840", "attention.decomposable = not self.fast_path", "compute_dtype=self.trt.float16", + "output_dtype=self.trt.float16", + 'tensor = self.cast(tensor, self.trt.float32, f"{name}.input_fp32")', + 'hidden = self.cast(hidden, self.trt.float16, "vit.residual_fp16")', + "compute_dtype=self.trt.float16 if self.fast_path else self.trt.float32", "config.builder_optimization_level = 3 if fast_path else 1", "config.avg_timing_iterations = 3", + "ElementWiseOperation.FLOOR_DIV", + "add_gather", + '"output.valid_fp16"', + "np.full((1,) * len(tuple(tensor.shape))", + "[[[0.5]]]", + '"output.raw_xy"', + '"output.raw_z"', + '"output.focal_samples_nchw"', + '"output.affine_depth_fp32"', + '"output.focal_samples_fp32"', + '"output.mask_sigmoid"', + '"output.mask_fp32"', ): assert required in source - for output in ("points", "mask", "metric_scale"): + assert '"output.points_nhwc"' not in source + assert '"output.points_remap"' not in source + assert '"output.valid.mask_finite"' not in source + assert '"output.valid_int8"' not in source + for output in ("affine_depth", "valid", "focal_samples", "metric_scale"): assert f'("{output}",' in source +def test_focal_sample_index_contract_covers_observed_shapes() -> None: + observed_sizes = ( + (608, 1080), + (612, 1080), + (1066, 1920), + (1076, 1920), + (1078, 1920), + (1080, 1840), + (1080, 1904), + (1080, 1906), + (1080, 1912), + (1080, 1918), + (1080, 1920), + (1264, 1080), + (1428, 1080), + (1440, 1080), + (1674, 1080), + (1904, 1080), + (1906, 1080), + (1912, 1080), + (1918, 1074), + (1918, 1080), + (1920, 1076), + (1920, 1078), + (1920, 1080), + (2688, 1508), + (3840, 2156), + (3840, 2160), + ) + + assert len(observed_sizes) == 26 + sample_size = model_module._FOCAL_RECOVERY_SIZE + assert sample_size == 64 + for width, height in observed_sizes: + for size in (width, height): + indices = tuple(index * size // sample_size for index in range(sample_size)) + assert indices[0] == 0 + assert indices[-1] == (sample_size - 1) * size // sample_size + assert all(0 <= index < size for index in indices) + assert all(left <= right for left, right in zip(indices, indices[1:])) + + +def test_slim_graph_retains_legacy_mask_rounding_and_ieee_finite_edges() -> None: + tiny_positive = np.nextafter(np.float16(0.0), np.float16(1.0)) + logits = np.asarray([-np.inf, -0.0, tiny_positive, np.inf, np.nan], dtype=np.float16) + with np.errstate(over="ignore", invalid="ignore"): + probabilities = np.asarray( + 1.0 / (1.0 + np.exp(-logits.astype(np.float32))), dtype=np.float16 + ) + legacy_selected = np.isfinite(probabilities) & (probabilities > np.float16(0.5)) + ordered_probability_selected = probabilities > np.float16(0.5) + logit_selected = logits > np.float16(0.0) + + np.testing.assert_array_equal(legacy_selected, np.asarray([False, False, False, True, False])) + np.testing.assert_array_equal(ordered_probability_selected, legacy_selected) + np.testing.assert_array_equal(logit_selected, np.asarray([False, False, True, True, False])) + assert probabilities[2] == np.float16(0.5) + + values = np.asarray( + [-np.finfo(np.float16).max, np.finfo(np.float16).max, -np.inf, np.inf, np.nan], + dtype=np.float16, + ) + with np.errstate(invalid="ignore"): + ordered_finite = np.abs(values) < np.float16(np.inf) + np.testing.assert_array_equal(ordered_finite, np.isfinite(values)) + + +def test_fp16_valid_output_uses_exact_zero_and_one_bit_patterns() -> None: + valid = np.asarray([False, True], dtype=np.bool_).astype(np.float16) + np.testing.assert_array_equal(valid.view(np.uint16), np.asarray([0x0000, 0x3C00], np.uint16)) + + +def test_slim_sample_gather_preserves_the_legacy_fp16_cast_boundary() -> None: + height, width = 67, 83 + affine_nchw = np.arange(3 * height * width, dtype=np.float32).reshape(1, 3, height, width) + affine_nchw = np.asarray(affine_nchw / 257.0, dtype=np.float16) + rows = np.asarray([index * height // 64 for index in range(64)]) + columns = np.asarray([index * width // 64 for index in range(64)]) + + legacy_nhwc = np.transpose(affine_nchw, (0, 2, 3, 1)).astype(np.float32) + legacy_samples = legacy_nhwc[:, rows, :, :][:, :, columns, :] + sampled_nchw = affine_nchw[:, :, rows, :][:, :, :, columns] + slim_samples = np.transpose(sampled_nchw, (0, 2, 3, 1)).astype(np.float32) + np.testing.assert_array_equal(slim_samples, legacy_samples) + + legacy_depth = legacy_nhwc[..., 2] + slim_depth = affine_nchw[:, 2, :, :].astype(np.float32) + np.testing.assert_array_equal(slim_depth, legacy_depth) + + def test_build_rejects_unknown_precision_and_wrong_checkpoint(tmp_path: Path) -> None: (tmp_path / "model.pt").write_bytes(b"wrong checkpoint") diff --git a/python/tensorrt_model_connect/families/moge/tests/test_fused_resample.py b/python/tensorrt_model_connect/families/moge/tests/test_fused_resample.py new file mode 100644 index 0000000000..4a8dae3993 --- /dev/null +++ b/python/tensorrt_model_connect/families/moge/tests/test_fused_resample.py @@ -0,0 +1,240 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only algebra and graph contracts for the exact level-3 resample fusion.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +from tensorrt_model_connect.families.moge import model as model_module + + +def _half_pixel_resize_x2(tensor: np.ndarray) -> np.ndarray: + batch, channels, height, width = tensor.shape + output = np.empty((batch, channels, 2 * height, 2 * width), dtype=tensor.dtype) + for output_y in range(2 * height): + source_y = (output_y + 0.5) * 0.5 - 0.5 + lower_y_raw = int(np.floor(source_y)) + fraction_y = source_y - lower_y_raw + lower_y = min(height - 1, max(0, lower_y_raw)) + upper_y = min(height - 1, max(0, lower_y_raw + 1)) + for output_x in range(2 * width): + source_x = (output_x + 0.5) * 0.5 - 0.5 + lower_x_raw = int(np.floor(source_x)) + fraction_x = source_x - lower_x_raw + lower_x = min(width - 1, max(0, lower_x_raw)) + upper_x = min(width - 1, max(0, lower_x_raw + 1)) + output[:, :, output_y, output_x] = ( + tensor[:, :, lower_y, lower_x] * (1.0 - fraction_y) * (1.0 - fraction_x) + + tensor[:, :, lower_y, upper_x] * (1.0 - fraction_y) * fraction_x + + tensor[:, :, upper_y, lower_x] * fraction_y * (1.0 - fraction_x) + + tensor[:, :, upper_y, upper_x] * fraction_y * fraction_x + ) + return output + + +def _replicate_conv3x3(tensor: np.ndarray, weight: np.ndarray, bias: np.ndarray) -> np.ndarray: + padded = np.pad(tensor, ((0, 0), (0, 0), (1, 1), (1, 1)), mode="edge") + output = np.empty( + (tensor.shape[0], weight.shape[0], tensor.shape[2], tensor.shape[3]), + dtype=tensor.dtype, + ) + for output_y in range(tensor.shape[2]): + for output_x in range(tensor.shape[3]): + patch = padded[:, :, output_y : output_y + 3, output_x : output_x + 3] + output[:, :, output_y, output_x] = np.einsum("nchw,ochw->no", patch, weight) + bias + return output + + +def _deconvolution_stride2_padding4( + tensor: np.ndarray, weight: np.ndarray, bias: np.ndarray +) -> np.ndarray: + full_height = (tensor.shape[2] - 1) * 2 + 6 + full_width = (tensor.shape[3] - 1) * 2 + 6 + full = np.zeros( + (tensor.shape[0], weight.shape[1], full_height, full_width), + dtype=tensor.dtype, + ) + for input_y in range(tensor.shape[2]): + for input_x in range(tensor.shape[3]): + contribution = np.einsum("ni,iohw->nohw", tensor[:, :, input_y, input_x], weight) + full[ + :, + :, + 2 * input_y : 2 * input_y + 6, + 2 * input_x : 2 * input_x + 6, + ] += contribution + output_height = (tensor.shape[2] - 1) * 2 - 8 + 6 + output_width = (tensor.shape[3] - 1) * 2 - 8 + 6 + return full[:, :, 4 : 4 + output_height, 4 : 4 + output_width] + bias[None, :, None, None] + + +def _reference(tensor: np.ndarray, weight: np.ndarray, bias: np.ndarray) -> np.ndarray: + return _replicate_conv3x3(_half_pixel_resize_x2(tensor), weight, bias) + + +def _fused(tensor: np.ndarray, weight: np.ndarray, bias: np.ndarray) -> np.ndarray: + padded = np.pad(tensor, ((0, 0), (0, 0), (1, 1), (1, 1)), mode="edge") + fused_weight = model_module._fuse_half_pixel_x2_conv_weight(weight) + return _deconvolution_stride2_padding4(padded, fused_weight, bias) + + +@pytest.mark.parametrize( + ("dtype", "tolerance"), + ((np.float64, 1.0e-11), (np.float32, 5.0e-5)), +) +@pytest.mark.parametrize("height,width", ((1, 1), (1, 3), (2, 1), (2, 2), (3, 5), (7, 4))) +def test_fused_resample_matches_half_pixel_replicate_reference( + dtype: Any, tolerance: float, height: int, width: int +) -> None: + random = np.random.default_rng(1000 + 10 * height + width) + tensor = random.standard_normal((1, 3, height, width)).astype(dtype) + weight = random.standard_normal((2, 3, 3, 3)).astype(dtype) + bias = random.standard_normal((2,)).astype(dtype) + + reference = _reference(tensor, weight, bias) + fused = _fused(tensor, weight, bias) + + assert reference.shape == fused.shape == (1, 2, 2 * height, 2 * width) + np.testing.assert_allclose(fused, reference, rtol=0.0, atol=tolerance) + + +@pytest.mark.parametrize("height,width", ((1, 1), (2, 3), (5, 4))) +def test_fused_resample_preserves_corner_edge_and_center_impulses(height: int, width: int) -> None: + random = np.random.default_rng(2000 + 10 * height + width) + weight = random.standard_normal((2, 2, 3, 3)).astype(np.float64) + bias = np.zeros((2,), dtype=np.float64) + positions = { + (0, 0), + (0, width - 1), + (height - 1, 0), + (height - 1, width - 1), + (height // 2, width // 2), + } + for input_y, input_x in positions: + tensor = np.zeros((1, 2, height, width), dtype=np.float64) + tensor[0, 0, input_y, input_x] = 1.0 + tensor[0, 1, input_y, input_x] = -0.5 + np.testing.assert_allclose( + _fused(tensor, weight, bias), + _reference(tensor, weight, bias), + rtol=0.0, + atol=1.0e-11, + ) + + +@dataclass +class _FakeTensor: + dtype: Any + shape: tuple[int, ...] + + +class _FakeLayer: + def __init__(self, output: _FakeTensor) -> None: + self.output = output + self.name = "" + self.stride_nd: tuple[int, ...] | None = None + self.padding_nd: tuple[int, ...] | None = None + + def get_output(self, index: int) -> _FakeTensor: + assert index == 0 + return self.output + + +class _FakeNetwork: + def __init__(self, dtype: Any) -> None: + self.dtype = dtype + self.calls: list[dict[str, Any]] = [] + + def add_deconvolution_nd( + self, + tensor: _FakeTensor, + output_channels: int, + kernel_size: tuple[int, int], + weight: np.ndarray, + bias: np.ndarray, + ) -> _FakeLayer: + height = (tensor.shape[2] - 1) * 2 + kernel_size[0] - 8 + width = (tensor.shape[3] - 1) * 2 + kernel_size[1] - 8 + layer = _FakeLayer(_FakeTensor(self.dtype, (1, output_channels, height, width))) + self.calls.append( + { + "tensor": tensor, + "output_channels": output_channels, + "kernel_size": kernel_size, + "weight": weight, + "bias": bias, + "layer": layer, + } + ) + return layer + + +class _FakeTrt: + float16 = "float16" + float32 = "float32" + + @staticmethod + def Weights(array: np.ndarray) -> np.ndarray: + return array + + +@pytest.mark.parametrize("height,width", ((1, 1), (3, 5), (32, 57), (57, 32))) +def test_fused_resample_graph_contract_is_dynamic_and_native( + monkeypatch: pytest.MonkeyPatch, height: int, width: int +) -> None: + network = _FakeNetwork(_FakeTrt.float32) + graph = model_module._NativeMogeGraph(_FakeTrt, network, {}, fast_path=True) + source_weight = np.arange(2 * 3 * 3 * 3, dtype=np.float32).reshape(2, 3, 3, 3) + source_bias = np.asarray((0.25, -0.5), dtype=np.float32) + arrays = {"resampler.weight": source_weight, "resampler.bias": source_bias} + monkeypatch.setattr(graph, "_array", lambda name, expected=None: arrays[name]) + pad_calls: list[tuple[int, str]] = [] + + def fake_pad(tensor: _FakeTensor, padding: int, name: str) -> _FakeTensor: + pad_calls.append((padding, name)) + return _FakeTensor( + tensor.dtype, + (tensor.shape[0], tensor.shape[1], tensor.shape[2] + 2, tensor.shape[3] + 2), + ) + + monkeypatch.setattr(graph, "replicate_pad", fake_pad) + tensor = _FakeTensor(_FakeTrt.float32, (1, 3, height, width)) + + output = graph.fused_half_pixel_resample( + tensor, + "resampler", + "level3", + compute_dtype=_FakeTrt.float32, + ) + + assert output.shape == (1, 2, 2 * height, 2 * width) + assert pad_calls == [(1, "level3.input_pad")] + assert len(network.calls) == 1 + call = network.calls[0] + assert call["output_channels"] == 2 + assert call["kernel_size"] == (6, 6) + assert call["layer"].stride_nd == (2, 2) + assert call["layer"].padding_nd == (4, 4) + np.testing.assert_array_equal( + call["weight"], model_module._fuse_half_pixel_x2_conv_weight(source_weight) + ) + np.testing.assert_array_equal(call["bias"], source_bias) + + +def test_level3_fast_path_uses_fused_native_deconvolution_only() -> None: + source = Path(model_module.__file__).read_text(encoding="utf-8") + assert "elif level == 3 and self.fast_path:" in source + assert 'f"{name}.fused_deconvolution"' in source + assert "self.network.add_deconvolution_nd" in source + assert "layer.stride_nd = (2, 2)" in source + assert "layer.padding_nd = (4, 4)" in source + for stack in ('"neck"', '"points_head"', '"mask_head"'): + assert stack in source + assert "add_plugin" not in source.lower() diff --git a/python/tensorrt_model_connect/families/moge/tests/test_selective_zero_padding.py b/python/tensorrt_model_connect/families/moge/tests/test_selective_zero_padding.py new file mode 100644 index 0000000000..3ccf8b0a7f --- /dev/null +++ b/python/tensorrt_model_connect/families/moge/tests/test_selective_zero_padding.py @@ -0,0 +1,210 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU-only contracts for the family-owned selective zero-padding lane.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +from tensorrt_model_connect.families.moge import model as model_module + + +EXPECTED_ZERO_PAD_SELECTION = frozenset( + { + "mask_head.res_blocks.1.0.layers.5", + "mask_head.res_blocks.2.0.layers.2", + "mask_head.res_blocks.2.0.layers.5", + "mask_head.res_blocks.3.0.layers.2", + "mask_head.res_blocks.3.0.layers.5", + "mask_head.resamplers.1.1", + "mask_head.resamplers.2.1", + "neck.res_blocks.2.0.layers.2", + "neck.res_blocks.2.0.layers.5", + "neck.res_blocks.2.1.layers.2", + "neck.res_blocks.2.1.layers.5", + "neck.res_blocks.3.0.layers.2", + "neck.res_blocks.3.0.layers.5", + "neck.res_blocks.3.1.layers.2", + "neck.res_blocks.3.1.layers.5", + "points_head.res_blocks.1.0.layers.2", + "points_head.res_blocks.1.0.layers.5", + "points_head.res_blocks.2.0.layers.2", + "points_head.res_blocks.2.0.layers.5", + "points_head.res_blocks.3.0.layers.2", + "points_head.res_blocks.3.0.layers.5", + "points_head.resamplers.0.1", + "points_head.resamplers.1.1", + "points_head.resamplers.2.1", + } +) +FUSED_LEVEL3_RESAMPLERS = frozenset( + { + "neck.resamplers.3.1", + "points_head.resamplers.3.1", + "mask_head.resamplers.3.1", + } +) + + +def _decoder_replicate_modules() -> frozenset[str]: + modules = set() + for prefix, block_counts in ( + ("neck", (0, 2, 2, 2, 0)), + ("points_head", (0, 1, 1, 1, 0)), + ("mask_head", (0, 1, 1, 1, 0)), + ): + modules.update(f"{prefix}.resamplers.{level}.1" for level in range(4)) + for level, count in enumerate(block_counts): + for block in range(count): + modules.add(f"{prefix}.res_blocks.{level}.{block}.layers.2") + modules.add(f"{prefix}.res_blocks.{level}.{block}.layers.5") + return frozenset(modules) + + +@dataclass +class _FakeTensor: + dtype: Any + shape: tuple[int, ...] + label: str + + +class _FakeLayer: + def __init__(self, output: _FakeTensor) -> None: + self.output = output + self.name = "" + self.stride_nd = None + self.padding_nd = None + + def get_output(self, index: int) -> _FakeTensor: + assert index == 0 + return self.output + + +class _FakeNetwork: + def __init__(self) -> None: + self.calls = [] + + def add_convolution_nd( + self, + tensor: _FakeTensor, + output_channels: int, + kernel_shape: tuple[int, int], + weight: np.ndarray, + bias: np.ndarray, + ) -> _FakeLayer: + layer = _FakeLayer(_FakeTensor(tensor.dtype, (1, output_channels, 8, 8), "output")) + self.calls.append( + { + "tensor": tensor, + "kernel_shape": kernel_shape, + "weight": weight, + "bias": bias, + "layer": layer, + } + ) + return layer + + +class _FakeTrt: + float16 = "float16" + float32 = "float32" + + @staticmethod + def Weights(array: np.ndarray) -> np.ndarray: + return array + + +def test_zero_pad_selection_is_exact_unique_and_topology_owned() -> None: + replicate_modules = _decoder_replicate_modules() + + assert len(replicate_modules) == 36 + assert model_module._ZERO_PAD_SELECTION == EXPECTED_ZERO_PAD_SELECTION + assert len(model_module._ZERO_PAD_SELECTION) == 24 + assert model_module._ZERO_PAD_SELECTION < replicate_modules + assert model_module._ZERO_PAD_SELECTION.isdisjoint(FUSED_LEVEL3_RESAMPLERS) + assert FUSED_LEVEL3_RESAMPLERS < replicate_modules + + +def test_only_selected_modules_use_native_zero_padding( + monkeypatch: pytest.MonkeyPatch, +) -> None: + network = _FakeNetwork() + graph = model_module._NativeMogeGraph(_FakeTrt, network, {}, fast_path=True) + weight = np.zeros((4, 4, 3, 3), dtype=np.float32) + bias = np.zeros((4,), dtype=np.float32) + monkeypatch.setattr( + graph, + "_array", + lambda name, expected=None: weight if name.endswith(".weight") else bias, + ) + monkeypatch.setattr(graph, "cast", lambda tensor, dtype, name: tensor) + pad_calls = [] + + def fake_replicate_pad(tensor: _FakeTensor, padding: int, name: str) -> _FakeTensor: + pad_calls.append((padding, name)) + return _FakeTensor(tensor.dtype, (1, 4, 10, 10), "replicate-padded") + + monkeypatch.setattr(graph, "replicate_pad", fake_replicate_pad) + + for module in sorted(_decoder_replicate_modules()): + pad_calls.clear() + network.calls.clear() + tensor = _FakeTensor(_FakeTrt.float32, (1, 4, 8, 8), "source") + graph.convolution( + tensor, + module, + f"graph.{module}", + replicate_padding=1, + compute_dtype=_FakeTrt.float32, + ) + assert len(network.calls) == 1 + call = network.calls[0] + if module in EXPECTED_ZERO_PAD_SELECTION: + assert pad_calls == [] + assert call["tensor"] is tensor + assert call["layer"].padding_nd == (1, 1) + else: + assert pad_calls == [(1, f"graph.{module}.pad")] + assert call["tensor"].label == "replicate-padded" + assert call["layer"].padding_nd == (0, 0) + + +def test_selected_module_rejects_an_unexpected_convolution_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + network = _FakeNetwork() + graph = model_module._NativeMogeGraph(_FakeTrt, network, {}, fast_path=True) + weight = np.zeros((4, 4, 3, 3), dtype=np.float32) + bias = np.zeros((4,), dtype=np.float32) + monkeypatch.setattr( + graph, + "_array", + lambda name, expected=None: weight if name.endswith(".weight") else bias, + ) + monkeypatch.setattr(graph, "cast", lambda tensor, dtype, name: tensor) + tensor = _FakeTensor(_FakeTrt.float32, (1, 4, 8, 8), "source") + + with pytest.raises(ValueError, match="must be stride-1 3x3"): + graph.convolution( + tensor, + next(iter(EXPECTED_ZERO_PAD_SELECTION)), + "invalid", + replicate_padding=0, + compute_dtype=_FakeTrt.float32, + ) + + +def test_selective_zero_patch_contains_no_linear_or_level4_fusion() -> None: + source = Path(model_module.__file__).read_text(encoding="utf-8") + for forbidden in ( + "_ENABLE_LEVEL12_LINEAR_FUSION", + "_ENABLE_LEVEL4_HEAD_FUSION", + "_compose_deconv2_replicate_conv3", + ): + assert forbidden not in source diff --git a/src/runtime/models/moge/pipeline.cpp b/src/runtime/models/moge/pipeline.cpp index ef754fcf51..1068f7f9fc 100644 --- a/src/runtime/models/moge/pipeline.cpp +++ b/src/runtime/models/moge/pipeline.cpp @@ -22,10 +22,9 @@ namespace { constexpr int32_t kFocalRecoverySize = 64; constexpr double kDenominatorEpsilon = 1.0e-9; constexpr int32_t kMinImageSize = 64; -constexpr int32_t kMaxImageSize = 2048; +constexpr int32_t kMaxImageSize = 4096; constexpr float kMinAspectRatio = 0.5F; constexpr float kMaxAspectRatio = 2.0F; -constexpr float kMaskThreshold = 0.5F; struct FocalSample { double u{0.0}; @@ -44,7 +43,26 @@ double normalized_coordinate(int32_t index, int32_t size, double span) { return span * (2.0 * index + 1.0 - size) / size; } -std::vector make_focal_samples(const float* points, const std::vector& mask, +int32_t nearest_sample_index(int32_t output_index, int32_t input_size) { + return std::min(input_size - 1, static_cast(static_cast(output_index) * + input_size / kFocalRecoverySize)); +} + +bool valid_focal_neighborhood(const uint16_t* valid, int32_t height, int32_t width, int32_t y, + int32_t x) { + if (y <= 0 || x <= 0 || y >= height - 1 || x >= width - 1) + return false; + for (int32_t neighbor_y = y - 1; neighbor_y <= y + 1; ++neighbor_y) { + for (int32_t neighbor_x = x - 1; neighbor_x <= x + 1; ++neighbor_x) { + const auto neighbor = static_cast(neighbor_y) * width + neighbor_x; + if (valid[neighbor] == 0) + return false; + } + } + return true; +} + +std::vector make_focal_samples(const float* sampled_points, const uint16_t* valid, int32_t height, int32_t width) { const double aspect = static_cast(width) / height; const double diagonal_factor = std::sqrt(1.0 + aspect * aspect); @@ -57,18 +75,18 @@ std::vector make_focal_samples(const float* points, const std::vect // floor(output_index * input_size / output_size). The official MoGe // recovery downsamples points, UVs, and the mask this way to 64x64. for (int32_t out_y = 0; out_y < kFocalRecoverySize; ++out_y) { - const int32_t y = std::min(height - 1, static_cast(static_cast(out_y) * - height / kFocalRecoverySize)); + const int32_t y = nearest_sample_index(out_y, height); for (int32_t out_x = 0; out_x < kFocalRecoverySize; ++out_x) { - const int32_t x = std::min(width - 1, static_cast(static_cast(out_x) * - width / kFocalRecoverySize)); + const int32_t x = nearest_sample_index(out_x, width); const auto pixel = static_cast(y) * width + x; - if (mask[pixel] == 0) + if (valid[pixel] == 0) continue; - const auto point = pixel * 3U; - const double px = points[point]; - const double py = points[point + 1U]; - const double pz = points[point + 2U]; + if (!valid_focal_neighborhood(valid, height, width, y, x)) + continue; + const auto sample = (static_cast(out_y) * kFocalRecoverySize + out_x) * 3U; + const double px = sampled_points[sample]; + const double py = sampled_points[sample + 1U]; + const double pz = sampled_points[sample + 2U]; if (!std::isfinite(px) || !std::isfinite(py) || !std::isfinite(pz)) continue; samples.push_back({normalized_coordinate(x, width, span_x), @@ -171,29 +189,24 @@ const float* require_float_output(const TensorMap& outputs, const char* name, return static_cast(tensor.data); } +const uint16_t* require_float16_output(const TensorMap& outputs, const char* name, + const std::vector& shape) { + const auto iterator = outputs.find(name); + if (iterator == outputs.end()) + throw std::runtime_error(std::string("MoGe engine did not return required output '") + + name + "'"); + const auto& tensor = iterator->second; + if (tensor.data == nullptr || tensor.dtype != DType::kFloat16 || tensor.shape != shape) { + throw std::runtime_error(std::string("MoGe output contract mismatch for '") + name + "'"); + } + return static_cast(tensor.data); +} + void validate_metric_scale(float metric_scale) { if (!std::isfinite(metric_scale) || metric_scale <= 0.0F) throw std::invalid_argument("MoGe metric scale must be finite and positive"); } -bool valid_raw_geometry_pixel(const float* affine_points, const float* mask_probabilities, - std::size_t pixel) { - const auto point = pixel * 3U; - return std::isfinite(mask_probabilities[pixel]) && mask_probabilities[pixel] > kMaskThreshold && - std::isfinite(affine_points[point]) && std::isfinite(affine_points[point + 1U]) && - std::isfinite(affine_points[point + 2U]); -} - -std::vector make_raw_mask(const float* affine_points, const float* mask_probabilities, - std::size_t area) { - std::vector raw_mask(area, 0); - for (std::size_t pixel = 0; pixel < area; ++pixel) { - raw_mask[pixel] = static_cast( - valid_raw_geometry_pixel(affine_points, mask_probabilities, pixel)); - } - return raw_mask; -} - struct MogeCalibration { double span_x{1.0}; double span_y{1.0}; @@ -243,8 +256,8 @@ void write_invalid_geometry(moge::GeometryResult& result, std::size_t pixel, flo result.points[point + 2U] = infinity; } -void populate_geometry_result(moge::GeometryResult& result, const float* affine_points, - const std::vector& raw_mask, const FocalShift& recovered, +void populate_geometry_result(moge::GeometryResult& result, const float* affine_depth, + const uint16_t* valid_pixels, const FocalShift& recovered, const MogeCalibration& calibration, float metric_scale) { const float infinity = std::numeric_limits::infinity(); for (int32_t y = 0; y < result.height; ++y) { @@ -253,8 +266,9 @@ void populate_geometry_result(moge::GeometryResult& result, const float* affine_ const auto pixel = static_cast(y) * result.width + x; const auto point = pixel * 3U; const double depth_unscaled = - static_cast(affine_points[point + 2U]) + recovered.shift; - const bool valid = valid_metric_depth(raw_mask[pixel], depth_unscaled); + static_cast(affine_depth[pixel]) + recovered.shift; + const bool valid = + valid_metric_depth(static_cast(valid_pixels[pixel] != 0), depth_unscaled); result.mask[pixel] = static_cast(valid); if (!valid) { write_invalid_geometry(result, pixel, infinity); @@ -281,7 +295,7 @@ bool supported_aspect_ratio(int32_t height, int32_t width) { } bool valid_rgb_value(float value) { - return std::isfinite(value) && value >= 0.0F && value <= 1.0F; + return value >= 0.0F && value <= 1.0F; } void validate_image_input(const float* pixels, int32_t height, int32_t width) { @@ -298,10 +312,9 @@ void validate_image_input(const float* pixels, int32_t height, int32_t width) { } } -std::optional recover_focal_shift(const float* affine_points, - const std::vector& mask, int32_t height, - int32_t width) { - const auto samples = make_focal_samples(affine_points, mask, height, width); +std::optional recover_focal_shift(const float* focal_samples, const uint16_t* valid, + int32_t height, int32_t width) { + const auto samples = make_focal_samples(focal_samples, valid, height, width); if (samples.size() < 2U) return std::nullopt; @@ -336,20 +349,17 @@ std::optional recover_focal_shift(const float* affine_points, return FocalShift{static_cast(current.focal), static_cast(shift)}; } -moge::GeometryResult postprocess_geometry(const float* affine_points, - const float* mask_probabilities, float metric_scale, +moge::GeometryResult postprocess_geometry(const float* affine_depth, const uint16_t* valid, + const float* focal_samples, float metric_scale, int32_t height, int32_t width) { const auto area = static_cast(height) * width; validate_metric_scale(metric_scale); - const auto raw_mask = make_raw_mask(affine_points, mask_probabilities, area); - - const auto recovered = recover_focal_shift(affine_points, raw_mask, height, width); + const auto recovered = recover_focal_shift(focal_samples, valid, height, width); if (!recovered) throw std::runtime_error("MoGe could not recover camera focal length and shift"); const auto calibration = make_calibration(*recovered, height, width); auto result = make_geometry_result(height, width, area, calibration); - populate_geometry_result(result, affine_points, raw_mask, *recovered, calibration, - metric_scale); + populate_geometry_result(result, affine_depth, valid, *recovered, calibration, metric_scale); return result; } @@ -372,10 +382,12 @@ moge::GeometryResult MogePipeline::estimate_geometry(const float* pixels, int32_ Tensor image{const_cast(pixels), {1, height, width, 3}, DType::kFloat32}; const auto outputs = model_->forward({{"image", image}}); - const auto* points = require_float_output(outputs, "points", {1, height, width, 3}); - const auto* mask = require_float_output(outputs, "mask", {1, height, width}); + const auto* affine_depth = require_float_output(outputs, "affine_depth", {1, height, width}); + const auto* valid = require_float16_output(outputs, "valid", {1, height, width}); + const auto* focal_samples = require_float_output( + outputs, "focal_samples", {1, kFocalRecoverySize, kFocalRecoverySize, 3}); const auto* scale = require_float_output(outputs, "metric_scale", {1}); - return postprocess_geometry(points, mask, scale[0], height, width); + return postprocess_geometry(affine_depth, valid, focal_samples, scale[0], height, width); } } // namespace trtmc diff --git a/src/runtime/models/moge/plugin.cpp b/src/runtime/models/moge/plugin.cpp index 5292155fc7..e41746137a 100644 --- a/src/runtime/models/moge/plugin.cpp +++ b/src/runtime/models/moge/plugin.cpp @@ -37,8 +37,7 @@ class MogePlugin final : public IPipelinePlugin { public: std::unique_ptr create(const PipelineContext& context) override { return std::make_unique(load_moge_engine(context), - context.bundle.info.model_id, - context.bundle.info.precision == "fp16"); + context.bundle.info.model_id, false); } }; diff --git a/tests/cpp/models/moge/test_moge_pipeline.cpp b/tests/cpp/models/moge/test_moge_pipeline.cpp index e715c1c15a..5b6f60f816 100644 --- a/tests/cpp/models/moge/test_moge_pipeline.cpp +++ b/tests/cpp/models/moge/test_moge_pipeline.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -50,13 +51,39 @@ std::vector affine_points(int32_t height, int32_t width, float focal, flo return points; } +std::vector affine_depth(const std::vector& points) { + std::vector depth(points.size() / 3U); + for (std::size_t pixel = 0; pixel < depth.size(); ++pixel) + depth[pixel] = points[pixel * 3U + 2U]; + return depth; +} + +std::vector focal_samples(const std::vector& points, int32_t height, int32_t width) { + constexpr int32_t sample_size = 64; + std::vector samples(static_cast(sample_size) * sample_size * 3U); + for (int32_t out_y = 0; out_y < sample_size; ++out_y) { + const int32_t y = static_cast(static_cast(out_y) * height / sample_size); + for (int32_t out_x = 0; out_x < sample_size; ++out_x) { + const int32_t x = + static_cast(static_cast(out_x) * width / sample_size); + const auto source = (static_cast(y) * width + x) * 3U; + const auto target = (static_cast(out_y) * sample_size + out_x) * 3U; + std::copy_n(points.data() + source, 3, samples.data() + target); + } + } + return samples; +} + class FakeMogeModule final : public trtmc::ITrtModule { public: - FakeMogeModule(int32_t height, int32_t width, bool invalidate_pixel = false) + FakeMogeModule(int32_t height, int32_t width, bool invalidate_pixel = false, + trtmc::DType valid_dtype = trtmc::DType::kFloat16) : height_(height), width_(width), points_(affine_points(height, width, 0.8F, 1.25F)), - mask_(static_cast(height) * width, 0.9F) { + depth_(affine_depth(points_)), samples_(focal_samples(points_, height, width)), + valid_(static_cast(height) * width, uint16_t{0x3C00}), + valid_dtype_(valid_dtype) { if (invalidate_pixel) - mask_[5] = 0.1F; + valid_[5] = 0; } trtmc::TensorMap forward(const trtmc::TensorMap& inputs) override { @@ -67,8 +94,9 @@ class FakeMogeModule final : public trtmc::ITrtModule { input_values.assign(values, values + image->second.numel()); } return { - {"points", {points_.data(), {1, height_, width_, 3}, trtmc::DType::kFloat32}}, - {"mask", {mask_.data(), {1, height_, width_}, trtmc::DType::kFloat32}}, + {"affine_depth", {depth_.data(), {1, height_, width_}, trtmc::DType::kFloat32}}, + {"valid", {valid_.data(), {1, height_, width_}, valid_dtype_}}, + {"focal_samples", {samples_.data(), {1, 64, 64, 3}, trtmc::DType::kFloat32}}, {"metric_scale", {scale_.data(), {1}, trtmc::DType::kFloat32}}, }; } @@ -84,16 +112,19 @@ class FakeMogeModule final : public trtmc::ITrtModule { std::vector output_info() const override { return {}; } bool has_input(const std::string& name) const override { return name == "image"; } bool has_output(const std::string& name) const override { - return name == "points" || name == "mask" || name == "metric_scale"; + return name == "affine_depth" || name == "valid" || name == "focal_samples" || + name == "metric_scale"; + } + trtmc::DType tensor_dtype(const std::string& name) const override { + return name == "valid" ? valid_dtype_ : trtmc::DType::kFloat32; } - trtmc::DType tensor_dtype(const std::string&) const override { return trtmc::DType::kFloat32; } std::vector tensor_shape(const std::string& name) const override { if (name == "image") return {1, height_, width_, 3}; - if (name == "points") - return {1, height_, width_, 3}; - if (name == "mask") + if (name == "affine_depth" || name == "valid") return {1, height_, width_}; + if (name == "focal_samples") + return {1, 64, 64, 3}; if (name == "metric_scale") return {1}; throw std::runtime_error("unknown fake tensor"); @@ -108,7 +139,23 @@ class FakeMogeModule final : public trtmc::ITrtModule { bool ok() const override { return true; } void keep_alive(std::shared_ptr) override {} - void invalidate_all() { std::fill(mask_.begin(), mask_.end(), 0.1F); } + void invalidate_all() { std::fill(valid_.begin(), valid_.end(), uint16_t{0}); } + void set_valid(int32_t y, int32_t x, bool value) { + valid_.at(static_cast(y) * width_ + x) = value ? uint16_t{0x3C00} : 0; + } + void set_valid_neighborhood(int32_t y, int32_t x) { + for (int32_t neighbor_y = y - 1; neighbor_y <= y + 1; ++neighbor_y) { + for (int32_t neighbor_x = x - 1; neighbor_x <= x + 1; ++neighbor_x) + set_valid(neighbor_y, neighbor_x, true); + } + } + void set_focal_sample(int32_t y, int32_t x, float px, float py, float pz) { + const auto sample = (static_cast(y) * 64U + x) * 3U; + samples_.at(sample) = px; + samples_.at(sample + 1U) = py; + samples_.at(sample + 2U) = pz; + } + uint16_t valid_bits(std::size_t pixel) const { return valid_.at(pixel); } std::vector input_shape; std::vector input_values; @@ -117,7 +164,10 @@ class FakeMogeModule final : public trtmc::ITrtModule { int32_t height_; int32_t width_; std::vector points_; - std::vector mask_; + std::vector depth_; + std::vector samples_; + std::vector valid_; + trtmc::DType valid_dtype_; std::vector scale_{2.0F}; }; @@ -163,11 +213,13 @@ void test_pipeline_recovers_metric_geometry_from_hwc_input() { check(dynamic_cast(&pipeline) != nullptr, "MoGe exposes its family-owned geometry contract"); check(std::string(pipeline.model_id()) == "moge-2-vitl", "MoGe model id"); + check(module_ptr->valid_bits(0) == 0x3C00, "MoGe valid=true uses FP16 one bits"); } void test_invalid_mask_materializes_infinity() { constexpr int32_t size = 64; auto module = std::make_unique(size, size, true); + auto* module_ptr = module.get(); trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); auto image = rgb_image(size, size); @@ -176,6 +228,24 @@ void test_invalid_mask_materializes_infinity() { check(result.mask[5] == 0, "MoGe invalid pixel mask cleared"); check(std::isinf(result.depth[5]), "MoGe invalid depth is infinity"); check(std::isinf(result.points[15]), "MoGe invalid point is infinity"); + check(module_ptr->valid_bits(5) == 0x0000, "MoGe valid=false uses FP16 zero bits"); + check(module_ptr->valid_bits(0) == 0x3C00, "MoGe retained valid pixel uses FP16 one bits"); +} + +void test_legacy_int8_valid_contract_is_rejected() { + constexpr int32_t size = 64; + auto module = std::make_unique(size, size, false, trtmc::DType::kInt8); + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + auto image = rgb_image(size, size); + + try { + (void)pipeline.estimate_geometry(image.data(), size, size); + check(false, "MoGe rejects legacy INT8 valid output"); + } catch (const std::runtime_error& error) { + check(std::string(error.what()).find("output contract mismatch for 'valid'") != + std::string::npos, + "MoGe legacy INT8 valid rejection is explicit"); + } } void test_focal_recovery_failure_is_reported() { @@ -194,6 +264,88 @@ void test_focal_recovery_failure_is_reported() { } } +void test_invalid_rgb_values_are_rejected() { + constexpr int32_t size = 64; + const std::vector invalid_values = { + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + -0.01F, + 1.01F, + }; + for (const float value : invalid_values) { + auto module = std::make_unique(size, size); + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + auto image = rgb_image(size, size); + image[0] = value; + try { + (void)pipeline.estimate_geometry(image.data(), size, size); + check(false, "MoGe rejects non-finite or out-of-range RGB input"); + } catch (const std::invalid_argument& error) { + check(std::string(error.what()).find("RGB input values") != std::string::npos, + "MoGe RGB rejection is explicit"); + } + } +} + +void test_focal_sampling_excludes_mapped_image_edges() { + constexpr int32_t size = 64; + auto module = std::make_unique(size, size); + for (int32_t index = 0; index < size; ++index) { + module->set_focal_sample(0, index, 1000.0F, -1000.0F, 0.1F); + module->set_focal_sample(size - 1, index, 1000.0F, -1000.0F, 0.1F); + module->set_focal_sample(index, 0, 1000.0F, -1000.0F, 0.1F); + module->set_focal_sample(index, size - 1, 1000.0F, -1000.0F, 0.1F); + } + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + auto image = rgb_image(size, size); + + const auto result = pipeline.estimate_geometry(image.data(), size, size); + + check(close(result.intrinsics[0], 0.5656854F), + "MoGe focal sampling excludes mapped image-edge points"); + check(close(result.intrinsics[4], 0.5656854F), + "MoGe image-edge exclusion preserves normalized fy"); +} + +void test_focal_sampling_excludes_invalid_three_by_three_neighborhood() { + constexpr int32_t size = 64; + constexpr int32_t sample_y = 20; + constexpr int32_t sample_x = 20; + auto module = std::make_unique(size, size); + module->set_focal_sample(sample_y, sample_x, 1000.0F, -1000.0F, 0.1F); + module->set_valid(sample_y, sample_x + 1, false); + auto* module_ptr = module.get(); + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + auto image = rgb_image(size, size); + + const auto result = pipeline.estimate_geometry(image.data(), size, size); + + check(module_ptr->valid_bits(static_cast(sample_y) * size + sample_x) == 0x3C00, + "MoGe focal center remains valid when its neighbor is invalid"); + check(close(result.intrinsics[0], 0.5656854F), + "MoGe focal sampling excludes a center with an invalid neighbor"); +} + +void test_focal_sampling_retains_complete_interior_neighborhoods() { + constexpr int32_t size = 64; + auto module = std::make_unique(size, size); + module->invalidate_all(); + for (int32_t y : {16, 32, 48}) { + for (int32_t x : {16, 32, 48}) + module->set_valid_neighborhood(y, x); + } + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + auto image = rgb_image(size, size); + + const auto result = pipeline.estimate_geometry(image.data(), size, size); + + check(close(result.intrinsics[0], 0.5656854F), + "MoGe focal sampling retains complete interior neighborhoods"); + check(result.mask[32U * size + 32U] == 1, + "MoGe retained interior focal center remains valid geometry"); +} + void test_fast_path_rejects_non_960x540_input() { constexpr int32_t size = 64; auto module = std::make_unique(size, size); @@ -214,7 +366,12 @@ void test_fast_path_rejects_non_960x540_input() { int main() { test_pipeline_recovers_metric_geometry_from_hwc_input(); test_invalid_mask_materializes_infinity(); + test_legacy_int8_valid_contract_is_rejected(); test_focal_recovery_failure_is_reported(); + test_invalid_rgb_values_are_rejected(); + test_focal_sampling_excludes_mapped_image_edges(); + test_focal_sampling_excludes_invalid_three_by_three_neighborhood(); + test_focal_sampling_retains_complete_interior_neighborhoods(); test_fast_path_rejects_non_960x540_input(); if (g_failures != 0) { std::cerr << g_failures << " MoGe pipeline test(s) failed\n"; From a74e801acf18a688d308f4c627fe175c171648f6 Mon Sep 17 00:00:00 2001 From: yifeif <277870278+yifeif-nv@users.noreply.github.com> Date: Wed, 2 Sep 2026 03:35:25 -0700 Subject: [PATCH 4/6] fix(moge): retain 960x540 dynamic input Lower the FP16 dynamic profile floor so the original 960x540 fast-path contract remains supported alongside the observed portrait, landscape, and 4K shapes. Remove the obsolete runtime fixed-shape flag now that the fast path is dynamic. The restored 960x540 engine build and unchanged geometry gates pass on L4. Signed-off-by: yifeif <277870278+yifeif-nv@users.noreply.github.com> --- .../families/moge/model.py | 2 +- .../families/moge/tests/test_family.py | 2 +- src/runtime/models/moge/pipeline.cpp | 10 ++-------- src/runtime/models/moge/pipeline.h | 4 +--- src/runtime/models/moge/plugin.cpp | 2 +- tests/cpp/models/moge/test_moge_pipeline.cpp | 16 ---------------- 6 files changed, 6 insertions(+), 30 deletions(-) diff --git a/python/tensorrt_model_connect/families/moge/model.py b/python/tensorrt_model_connect/families/moge/model.py index dc3e50878d..41706e744d 100644 --- a/python/tensorrt_model_connect/families/moge/model.py +++ b/python/tensorrt_model_connect/families/moge/model.py @@ -33,7 +33,7 @@ _MIN_IMAGE_SIZE = 64 _OPT_IMAGE_SIZE = 518 _MAX_IMAGE_SIZE = 2048 -_FAST_MIN_IMAGE_HEIGHT = 1074 +_FAST_MIN_IMAGE_HEIGHT = 540 _FAST_MIN_IMAGE_WIDTH = 608 _FAST_OPT_IMAGE_HEIGHT = 1080 _FAST_OPT_IMAGE_WIDTH = 1920 diff --git a/python/tensorrt_model_connect/families/moge/tests/test_family.py b/python/tensorrt_model_connect/families/moge/tests/test_family.py index 9578b32b03..3f3bae10e8 100644 --- a/python/tensorrt_model_connect/families/moge/tests/test_family.py +++ b/python/tensorrt_model_connect/families/moge/tests/test_family.py @@ -129,7 +129,7 @@ def test_production_builder_is_fixed_and_tensor_rt_native_only() -> None: "first_transpose=(0, 3, 1, 2)", "_NUM_TOKENS = 1800", "_FOCAL_RECOVERY_SIZE = 64", - "_FAST_MIN_IMAGE_HEIGHT = 1074", + "_FAST_MIN_IMAGE_HEIGHT = 540", "_FAST_MIN_IMAGE_WIDTH = 608", "_FAST_OPT_IMAGE_HEIGHT = 1080", "_FAST_OPT_IMAGE_WIDTH = 1920", diff --git a/src/runtime/models/moge/pipeline.cpp b/src/runtime/models/moge/pipeline.cpp index 1068f7f9fc..7335fcf06b 100644 --- a/src/runtime/models/moge/pipeline.cpp +++ b/src/runtime/models/moge/pipeline.cpp @@ -365,10 +365,8 @@ moge::GeometryResult postprocess_geometry(const float* affine_depth, const uint1 } // namespace -MogePipeline::MogePipeline(std::unique_ptr model, std::string model_id, - bool fixed_fast_shape) - : model_(std::move(model)), model_id_(std::move(model_id)), - fixed_fast_shape_(fixed_fast_shape) { +MogePipeline::MogePipeline(std::unique_ptr model, std::string model_id) + : model_(std::move(model)), model_id_(std::move(model_id)) { if (!model_ || !model_->ok()) throw std::runtime_error("MogePipeline: invalid model"); } @@ -376,10 +374,6 @@ MogePipeline::MogePipeline(std::unique_ptr model, std::string model_ moge::GeometryResult MogePipeline::estimate_geometry(const float* pixels, int32_t height, int32_t width) { validate_image_input(pixels, height, width); - if (fixed_fast_shape_ && (height != 540 || width != 960)) { - throw std::invalid_argument("MoGe FP16 fast path requires a 960x540 image"); - } - Tensor image{const_cast(pixels), {1, height, width, 3}, DType::kFloat32}; const auto outputs = model_->forward({{"image", image}}); const auto* affine_depth = require_float_output(outputs, "affine_depth", {1, height, width}); diff --git a/src/runtime/models/moge/pipeline.h b/src/runtime/models/moge/pipeline.h index eefbf53848..3046719657 100644 --- a/src/runtime/models/moge/pipeline.h +++ b/src/runtime/models/moge/pipeline.h @@ -17,8 +17,7 @@ namespace trtmc { class MogePipeline final : public IPipeline, public moge::IGeometryEstimator { public: - MogePipeline(std::unique_ptr model, std::string model_id, - bool fixed_fast_shape = false); + MogePipeline(std::unique_ptr model, std::string model_id); moge::GeometryResult estimate_geometry(const float* pixels, int32_t height, int32_t width) override; @@ -29,7 +28,6 @@ class MogePipeline final : public IPipeline, public moge::IGeometryEstimator { private: std::unique_ptr model_; std::string model_id_; - bool fixed_fast_shape_{false}; }; } // namespace trtmc diff --git a/src/runtime/models/moge/plugin.cpp b/src/runtime/models/moge/plugin.cpp index e41746137a..15fa9788d5 100644 --- a/src/runtime/models/moge/plugin.cpp +++ b/src/runtime/models/moge/plugin.cpp @@ -37,7 +37,7 @@ class MogePlugin final : public IPipelinePlugin { public: std::unique_ptr create(const PipelineContext& context) override { return std::make_unique(load_moge_engine(context), - context.bundle.info.model_id, false); + context.bundle.info.model_id); } }; diff --git a/tests/cpp/models/moge/test_moge_pipeline.cpp b/tests/cpp/models/moge/test_moge_pipeline.cpp index 5b6f60f816..5e64846080 100644 --- a/tests/cpp/models/moge/test_moge_pipeline.cpp +++ b/tests/cpp/models/moge/test_moge_pipeline.cpp @@ -346,21 +346,6 @@ void test_focal_sampling_retains_complete_interior_neighborhoods() { "MoGe retained interior focal center remains valid geometry"); } -void test_fast_path_rejects_non_960x540_input() { - constexpr int32_t size = 64; - auto module = std::make_unique(size, size); - trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl", true); - auto image = rgb_image(size, size); - - try { - (void)pipeline.estimate_geometry(image.data(), size, size); - check(false, "MoGe FP16 fast path rejects another shape"); - } catch (const std::invalid_argument& error) { - check(std::string(error.what()).find("960x540") != std::string::npos, - "MoGe FP16 fast-path shape error is explicit"); - } -} - } // namespace int main() { @@ -372,7 +357,6 @@ int main() { test_focal_sampling_excludes_mapped_image_edges(); test_focal_sampling_excludes_invalid_three_by_three_neighborhood(); test_focal_sampling_retains_complete_interior_neighborhoods(); - test_fast_path_rejects_non_960x540_input(); if (g_failures != 0) { std::cerr << g_failures << " MoGe pipeline test(s) failed\n"; return 1; From 98c8c559194aa594a28c958f28516f808eda2ef9 Mon Sep 17 00:00:00 2001 From: yifeif <277870278+yifeif-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:00:33 -0700 Subject: [PATCH 5/6] fix(moge): enforce engine profile bounds Read the active TensorRT image profile when constructing the MoGe pipeline and reject dimensions outside its min/max bounds before forward. Preserve the existing aspect-ratio and RGB validation, and cover active-profile, boundary, and malformed-profile behavior. Signed-off-by: yifeif <277870278+yifeif-nv@users.noreply.github.com> --- src/runtime/models/moge/pipeline.cpp | 47 ++++++++++-- src/runtime/models/moge/pipeline.h | 4 + tests/cpp/models/moge/test_moge_pipeline.cpp | 81 ++++++++++++++++++-- 3 files changed, 118 insertions(+), 14 deletions(-) diff --git a/src/runtime/models/moge/pipeline.cpp b/src/runtime/models/moge/pipeline.cpp index 7335fcf06b..e2f2ff61f7 100644 --- a/src/runtime/models/moge/pipeline.cpp +++ b/src/runtime/models/moge/pipeline.cpp @@ -21,11 +21,34 @@ namespace { constexpr int32_t kFocalRecoverySize = 64; constexpr double kDenominatorEpsilon = 1.0e-9; -constexpr int32_t kMinImageSize = 64; -constexpr int32_t kMaxImageSize = 4096; constexpr float kMinAspectRatio = 0.5F; constexpr float kMaxAspectRatio = 2.0F; +struct ImageProfileBounds { + int32_t min_height{0}; + int32_t min_width{0}; + int32_t max_height{0}; + int32_t max_width{0}; +}; + +bool valid_image_profile_shape(const std::vector& shape) { + constexpr auto kInt32Max = std::numeric_limits::max(); + return shape.size() == 4U && shape[0] == 1 && shape[1] > 0 && shape[1] <= kInt32Max && + shape[2] > 0 && shape[2] <= kInt32Max && shape[3] == 3; +} + +ImageProfileBounds image_profile_bounds(const ITrtModule& model) { + const auto profile = model.profile_idx(); + const auto minimum = model.input_profile_shape("image", profile, ProfileShapeSelector::kMin); + const auto maximum = model.input_profile_shape("image", profile, ProfileShapeSelector::kMax); + if (!valid_image_profile_shape(minimum) || !valid_image_profile_shape(maximum) || + minimum[1] > maximum[1] || minimum[2] > maximum[2]) { + throw std::runtime_error("MogePipeline: invalid TensorRT image profile"); + } + return {static_cast(minimum[1]), static_cast(minimum[2]), + static_cast(maximum[1]), static_cast(maximum[2])}; +} + struct FocalSample { double u{0.0}; double v{0.0}; @@ -284,9 +307,9 @@ void populate_geometry_result(moge::GeometryResult& result, const float* affine_ } } -bool supported_image_size(int32_t height, int32_t width) { - return height >= kMinImageSize && width >= kMinImageSize && height <= kMaxImageSize && - width <= kMaxImageSize; +bool supported_image_size(int32_t height, int32_t width, const ImageProfileBounds& profile) { + return height >= profile.min_height && width >= profile.min_width && + height <= profile.max_height && width <= profile.max_width; } bool supported_aspect_ratio(int32_t height, int32_t width) { @@ -298,10 +321,11 @@ bool valid_rgb_value(float value) { return value >= 0.0F && value <= 1.0F; } -void validate_image_input(const float* pixels, int32_t height, int32_t width) { +void validate_image_input(const float* pixels, int32_t height, int32_t width, + const ImageProfileBounds& profile) { if (pixels == nullptr) throw std::invalid_argument("MoGe image pointer is null"); - if (!supported_image_size(height, width)) + if (!supported_image_size(height, width, profile)) throw std::invalid_argument("MoGe image dimensions are outside the bundle profile"); if (!supported_aspect_ratio(height, width)) throw std::invalid_argument("MoGe image aspect ratio is outside the supported range"); @@ -369,11 +393,18 @@ MogePipeline::MogePipeline(std::unique_ptr model, std::string model_ : model_(std::move(model)), model_id_(std::move(model_id)) { if (!model_ || !model_->ok()) throw std::runtime_error("MogePipeline: invalid model"); + const auto profile = image_profile_bounds(*model_); + min_image_height_ = profile.min_height; + min_image_width_ = profile.min_width; + max_image_height_ = profile.max_height; + max_image_width_ = profile.max_width; } moge::GeometryResult MogePipeline::estimate_geometry(const float* pixels, int32_t height, int32_t width) { - validate_image_input(pixels, height, width); + validate_image_input( + pixels, height, width, + {min_image_height_, min_image_width_, max_image_height_, max_image_width_}); Tensor image{const_cast(pixels), {1, height, width, 3}, DType::kFloat32}; const auto outputs = model_->forward({{"image", image}}); const auto* affine_depth = require_float_output(outputs, "affine_depth", {1, height, width}); diff --git a/src/runtime/models/moge/pipeline.h b/src/runtime/models/moge/pipeline.h index 3046719657..dc4d97c3e9 100644 --- a/src/runtime/models/moge/pipeline.h +++ b/src/runtime/models/moge/pipeline.h @@ -28,6 +28,10 @@ class MogePipeline final : public IPipeline, public moge::IGeometryEstimator { private: std::unique_ptr model_; std::string model_id_; + int32_t min_image_height_{0}; + int32_t min_image_width_{0}; + int32_t max_image_height_{0}; + int32_t max_image_width_{0}; }; } // namespace trtmc diff --git a/tests/cpp/models/moge/test_moge_pipeline.cpp b/tests/cpp/models/moge/test_moge_pipeline.cpp index 5e64846080..a1ae6220df 100644 --- a/tests/cpp/models/moge/test_moge_pipeline.cpp +++ b/tests/cpp/models/moge/test_moge_pipeline.cpp @@ -78,8 +78,10 @@ class FakeMogeModule final : public trtmc::ITrtModule { public: FakeMogeModule(int32_t height, int32_t width, bool invalidate_pixel = false, trtmc::DType valid_dtype = trtmc::DType::kFloat16) - : height_(height), width_(width), points_(affine_points(height, width, 0.8F, 1.25F)), - depth_(affine_depth(points_)), samples_(focal_samples(points_, height, width)), + : height_(height), width_(width), min_height_(height), min_width_(width), + max_height_(height), max_width_(width), + points_(affine_points(height, width, 0.8F, 1.25F)), depth_(affine_depth(points_)), + samples_(focal_samples(points_, height, width)), valid_(static_cast(height) * width, uint16_t{0x3C00}), valid_dtype_(valid_dtype) { if (invalidate_pixel) @@ -87,6 +89,7 @@ class FakeMogeModule final : public trtmc::ITrtModule { } trtmc::TensorMap forward(const trtmc::TensorMap& inputs) override { + ++forward_count_; const auto image = inputs.find("image"); if (image != inputs.end()) { input_shape = image->second.shape; @@ -107,7 +110,7 @@ class FakeMogeModule final : public trtmc::ITrtModule { cudaStream_t stream() const override { return nullptr; } void enable_cuda_graph() override {} bool cuda_graph_active() const override { return false; } - int32_t profile_idx() const override { return 0; } + int32_t profile_idx() const override { return profile_index_; } std::vector input_info() const override { return {}; } std::vector output_info() const override { return {}; } bool has_input(const std::string& name) const override { return name == "image"; } @@ -129,11 +132,17 @@ class FakeMogeModule final : public trtmc::ITrtModule { return {1}; throw std::runtime_error("unknown fake tensor"); } - std::vector input_profile_shape(const std::string&, int32_t, - trtmc::ProfileShapeSelector) const override { + std::vector input_profile_shape(const std::string&, int32_t profile_index, + trtmc::ProfileShapeSelector selector) const override { + last_queried_profile_index_ = profile_index; + ++profile_query_count_; + if (selector == trtmc::ProfileShapeSelector::kMin) + return {1, min_height_, min_width_, 3}; + if (selector == trtmc::ProfileShapeSelector::kMax) + return {1, max_height_, max_width_, 3}; return {1, height_, width_, 3}; } - int32_t optimization_profile_count() const override { return 1; } + int32_t optimization_profile_count() const override { return profile_index_ + 1; } void* device_ptr(const std::string&) const override { return nullptr; } void bind_external(const std::string&, void*) override {} bool ok() const override { return true; } @@ -156,6 +165,17 @@ class FakeMogeModule final : public trtmc::ITrtModule { samples_.at(sample + 2U) = pz; } uint16_t valid_bits(std::size_t pixel) const { return valid_.at(pixel); } + int32_t forward_count() const { return forward_count_; } + int32_t last_queried_profile_index() const { return last_queried_profile_index_; } + int32_t profile_query_count() const { return profile_query_count_; } + void set_profile_index(int32_t profile_index) { profile_index_ = profile_index; } + void set_profile_bounds(int32_t min_height, int32_t min_width, int32_t max_height, + int32_t max_width) { + min_height_ = min_height; + min_width_ = min_width; + max_height_ = max_height; + max_width_ = max_width; + } std::vector input_shape; std::vector input_values; @@ -163,6 +183,14 @@ class FakeMogeModule final : public trtmc::ITrtModule { private: int32_t height_; int32_t width_; + int32_t min_height_; + int32_t min_width_; + int32_t max_height_; + int32_t max_width_; + int32_t forward_count_{0}; + int32_t profile_index_{0}; + mutable int32_t last_queried_profile_index_{-1}; + mutable int32_t profile_query_count_{0}; std::vector points_; std::vector depth_; std::vector samples_; @@ -288,6 +316,45 @@ void test_invalid_rgb_values_are_rejected() { } } +void test_loaded_engine_profile_bounds_are_enforced_before_forward() { + constexpr int32_t height = 64; + constexpr int32_t width = 80; + auto module = std::make_unique(height, width); + module->set_profile_index(2); + module->set_profile_bounds(height, width, 128, 160); + auto* module_ptr = module.get(); + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + check(module_ptr->last_queried_profile_index() == 2 && module_ptr->profile_query_count() == 2, + "MoGe reads min and max bounds from the active engine profile"); + + for (const auto& [input_height, input_width] : + {std::pair{height - 1, width}, std::pair{height, width - 1}, std::pair{129, width}, + std::pair{height, 161}}) { + auto image = rgb_image(input_height, input_width); + try { + (void)pipeline.estimate_geometry(image.data(), input_height, input_width); + check(false, "MoGe rejects dimensions outside the loaded engine profile"); + } catch (const std::invalid_argument& error) { + check(std::string(error.what()).find("outside the bundle profile") != std::string::npos, + "MoGe profile-bound rejection is explicit"); + } + } + check(module_ptr->forward_count() == 0, + "MoGe rejects profile-incompatible dimensions before TensorRT forward"); +} + +void test_invalid_loaded_engine_profile_is_rejected() { + auto module = std::make_unique(64, 80); + module->set_profile_bounds(128, 160, 64, 80); + try { + trtmc::MogePipeline pipeline(std::move(module), "moge-2-vitl"); + check(false, "MoGe rejects an invalid loaded engine profile"); + } catch (const std::runtime_error& error) { + check(std::string(error.what()).find("invalid TensorRT image profile") != std::string::npos, + "MoGe invalid engine profile rejection is explicit"); + } +} + void test_focal_sampling_excludes_mapped_image_edges() { constexpr int32_t size = 64; auto module = std::make_unique(size, size); @@ -354,6 +421,8 @@ int main() { test_legacy_int8_valid_contract_is_rejected(); test_focal_recovery_failure_is_reported(); test_invalid_rgb_values_are_rejected(); + test_loaded_engine_profile_bounds_are_enforced_before_forward(); + test_invalid_loaded_engine_profile_is_rejected(); test_focal_sampling_excludes_mapped_image_edges(); test_focal_sampling_excludes_invalid_three_by_three_neighborhood(); test_focal_sampling_retains_complete_interior_neighborhoods(); From 37148e850865961b8bb6a3c4f4ad2ad68505b21d Mon Sep 17 00:00:00 2001 From: yifeif <277870278+yifeif-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 17:12:29 -0700 Subject: [PATCH 6/6] fix(moge): keep FP32 attention decomposed Set the native FP32 builder to optimization level 0 so TensorRT retains the decomposed dynamic attention graph. Keep the FP16 fast path at level 3 and lock the family-owned policy in its focused contract test. Signed-off-by: yifeif <277870278+yifeif-nv@users.noreply.github.com> --- python/tensorrt_model_connect/families/moge/model.py | 7 +++---- .../families/moge/tests/test_family.py | 2 +- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/python/tensorrt_model_connect/families/moge/model.py b/python/tensorrt_model_connect/families/moge/model.py index 41706e744d..12b2d535cd 100644 --- a/python/tensorrt_model_connect/families/moge/model.py +++ b/python/tensorrt_model_connect/families/moge/model.py @@ -1373,10 +1373,9 @@ def _build_native_engine( raise RuntimeError("Failed to configure the MoGe dynamic image profile") config.add_optimization_profile(profile) if hasattr(config, "builder_optimization_level"): - # Level 3 enables the static FP16 fused-attention fast path. The broad - # dynamic FP32 graph stays at level 1 because level 3 absorbs its - # decomposable attention into an unsupported dynamic-BMM ForeignNode. - config.builder_optimization_level = 3 if fast_path else 1 + # Level 3 enables the FP16 fused-attention fast path. Level 0 keeps the + # broad dynamic FP32 attention graph decomposed for reliable builds. + config.builder_optimization_level = 3 if fast_path else 0 if hasattr(config, "avg_timing_iterations"): config.avg_timing_iterations = 3 if hasattr(config, "max_aux_streams"): diff --git a/python/tensorrt_model_connect/families/moge/tests/test_family.py b/python/tensorrt_model_connect/families/moge/tests/test_family.py index 3f3bae10e8..f8450be7a7 100644 --- a/python/tensorrt_model_connect/families/moge/tests/test_family.py +++ b/python/tensorrt_model_connect/families/moge/tests/test_family.py @@ -141,7 +141,7 @@ def test_production_builder_is_fixed_and_tensor_rt_native_only() -> None: 'tensor = self.cast(tensor, self.trt.float32, f"{name}.input_fp32")', 'hidden = self.cast(hidden, self.trt.float16, "vit.residual_fp16")', "compute_dtype=self.trt.float16 if self.fast_path else self.trt.float32", - "config.builder_optimization_level = 3 if fast_path else 1", + "config.builder_optimization_level = 3 if fast_path else 0", "config.avg_timing_iterations = 3", "ElementWiseOperation.FLOOR_DIV", "add_gather",