From dd6d726b1c19a9a306ff068ea2bb8a6448a786f7 Mon Sep 17 00:00:00 2001 From: Naren Dasan Date: Fri, 28 Aug 2026 01:16:21 +0000 Subject: [PATCH 1/2] fix(dynamo): normalize scalar values at engine boundaries Materialize scalar partition inputs as zero-dimensional CUDA tensors before calling the Tensor[] execute_engine schema, and restore scalar engine outputs with aten.item for their original consumers. The meta kernel now skips scalar input metadata and selects output placement from the first tensor input. Fixes #4614 Fixes #4616 Fixes #4617 Fixes #4620 Tests: - pytest tests/py/dynamo/models/test_exporter_inlining.py -n0 (7 passed) - pytest tests/py/dynamo/models/test_meta_kernel_shape_inference.py -n0 (12 passed) - original issue reproductions for #4614, #4616, #4617, and #4620 - #4620 AOTInductor package creation succeeded --- py/torch_tensorrt/dynamo/_exporter.py | 98 ++++++++++++++++--- .../dynamo/runtime/_TorchTensorRTModule.py | 9 +- .../runtime/meta_ops/register_meta_ops.py | 15 ++- .../dynamo/models/test_exporter_inlining.py | 93 ++++++++++++++++++ .../test_meta_kernel_shape_inference.py | 52 ++++++++++ 5 files changed, 249 insertions(+), 18 deletions(-) diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index c6c992d720..74c819fe73 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -888,6 +888,20 @@ def inline_trt_modules( """ Replace TRT submodules with trt engine nodes. """ + fake_mode = detect_fake_mode( + tuple( + node.meta["val"] + for node in gm.graph.nodes + if node.op == "placeholder" and "val" in node.meta + ) + ) + + def scalar_tensor_meta(dtype: torch.dtype, device: torch.device) -> torch.Tensor: + if fake_mode is None: + return torch.empty((), dtype=dtype, device="meta") + with fake_mode: + return torch.empty((), dtype=dtype, device=device) + for name, _ in gm.named_children(): if "_run_on_acc" not in name: continue @@ -904,16 +918,52 @@ def inline_trt_modules( raise ValueError( f"trt_module_node: {trt_module_node.name} does not have the metadata which should be set during dynamo compile_module step." ) - num_outputs = len(trt_module_node.meta["val"]) + + shape_info = trt_module.symbolic_shape_expressions or {} + input_info = shape_info.get("inputs", []) + output_info = shape_info.get("outputs", []) + original_output_vals = trt_module_node.meta["val"] + if not isinstance(original_output_vals, (tuple, list)): + original_output_vals = [original_output_vals] + num_outputs = len(original_output_vals) + + engine_output_vals = [] + for index, value in enumerate(original_output_vals): + if index < len(output_info) and output_info[index].get("is_scalar"): + engine_output_vals.append( + scalar_tensor_meta( + output_info[index]["dtype"], trt_module.target_device + ) + ) + else: + engine_output_vals.append(value) + # Insert a call_function node to perform inference on TRT engine with gm.graph.inserting_before(trt_module_node): + engine_inputs = list(trt_module_node.args) + for index, info in enumerate(input_info): + if index >= len(engine_inputs) or not info.get("is_scalar"): + continue + scalar_input = gm.graph.call_function( + torch.ops.aten.scalar_tensor.default, + (engine_inputs[index],), + { + "dtype": info["dtype"], + "device": trt_module.target_device, + }, + ) + scalar_input.meta["val"] = scalar_tensor_meta( + info["dtype"], trt_module.target_device + ) + engine_inputs[index] = scalar_input + if cross_compile_module: engine_info = trt_module._pack_engine_info() engine_bytes = engine_info[ENGINE_IDX] engine_info[ENGINE_IDX] = base64.b64encode(engine_bytes).decode("utf-8") trt_node = gm.graph.call_function( torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default, - (trt_module_node.args, *engine_info), + (tuple(engine_inputs), *engine_info), ) else: engine_name = f"{name}_engine" @@ -922,28 +972,50 @@ def inline_trt_modules( trt_node = gm.graph.call_function( torch.ops.tensorrt.execute_engine.default, - (trt_module_node.args, engine_node), + (tuple(engine_inputs), engine_node), ) engine_node.meta["val"] = CustomObjArgument( name=engine_node.name, class_fqn="" ) assert num_outputs > 0 - trt_node.meta["val"] = trt_module_node.meta["val"] + trt_node.meta["val"] = engine_output_vals + + def restore_scalar_output( + getitem_node: torch.fx.Node, output_index: int + ) -> torch.fx.Node: + getitem_node.meta["val"] = engine_output_vals[output_index] + if output_index >= len(output_info) or not output_info[output_index].get( + "is_scalar" + ): + return getitem_node + + with gm.graph.inserting_after(getitem_node): + scalar_output = gm.graph.call_function( + torch.ops.aten.item.default, (getitem_node,) + ) + scalar_output.meta["val"] = original_output_vals[output_index] + getitem_node.replace_all_uses_with( + scalar_output, delete_user_cb=lambda user: user is not scalar_output + ) + return scalar_output if num_outputs == 1: - # Insert getitem nodes as outputs (for export serialization to work) + # Insert a getitem because execute_engine always returns Tensor[]. with gm.graph.inserting_after(trt_node): getitem_output = gm.graph.call_function(operator.getitem, (trt_node, 0)) - getitem_output.meta["val"] = trt_node.meta["val"] - trt_module_node.replace_all_uses_with(getitem_output) + replacement = restore_scalar_output(getitem_output, 0) + trt_module_node.replace_all_uses_with(replacement) else: - # Multiple outputs case: - # Replace uses of submodule with the trt_node. - # getitem nodes are already added inherently by the partitioner + # Multiple-output partitioner graphs already contain getitem users. trt_module_node.replace_all_uses_with(trt_node) - getitem_nodes = trt_node.users - for idx, getitem_node in enumerate(getitem_nodes): - getitem_node.meta["val"] = trt_node.meta["val"][idx] + for getitem_node in list(trt_node.users): + if ( + getitem_node.target is not operator.getitem + or len(getitem_node.args) < 2 + or not isinstance(getitem_node.args[1], int) + ): + continue + restore_scalar_output(getitem_node, getitem_node.args[1]) # Expose the engine's aliased (KV-cache) outputs as graph-level buffer # mutations so the ExecuTorch path sees a real mutable buffer instead of diff --git a/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py b/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py index b52bb8d360..6c2f8c7e9e 100644 --- a/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py +++ b/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py @@ -711,7 +711,7 @@ def pre_allocated_outputs(self) -> Any: def set_use_output_allocator(self, enable: bool) -> None: self.get_engine().use_output_allocator_outputs = enable - def forward(self, *inputs: Any) -> torch.Tensor | Tuple[torch.Tensor, ...]: + def forward(self, *inputs: Any) -> Any: """Run the TensorRT engine on GPU tensors (non-tensor args are cast to CUDA tensors). Note: callers are responsible for ensuring the engine has been set up; @@ -745,7 +745,7 @@ def forward(self, *inputs: Any) -> torch.Tensor | Tuple[torch.Tensor, ...]: else: input_tensors.append(torch.tensor(i).cuda()) - outputs: List[torch.Tensor] = torch.ops.tensorrt.execute_engine( + outputs: List[Any] = torch.ops.tensorrt.execute_engine( list(input_tensors), self.engine ) @@ -760,6 +760,11 @@ def forward(self, *inputs: Any) -> torch.Tensor | Tuple[torch.Tensor, ...]: if n < len(outputs): outputs = outputs[:n] + output_info = (self.symbolic_shape_expressions or {}).get("outputs", []) + for index, info in enumerate(output_info): + if index < len(outputs) and info.get("is_scalar"): + outputs[index] = outputs[index].item() + if len(outputs) == 1: return outputs[0] diff --git a/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py b/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py index fcb6f21eb8..294e8ab72a 100644 --- a/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py +++ b/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py @@ -10,7 +10,7 @@ def _apply_symbolic_shape_expressions( - inputs: List[torch.Tensor], shape_info: Dict[str, List[Dict[str, Any]]] + inputs: List[Any], shape_info: Dict[str, List[Dict[str, Any]]] ) -> List[torch.Tensor]: """ Apply symbolic shape expressions to create output fake tensors. @@ -35,6 +35,13 @@ def _apply_symbolic_shape_expressions( input_info = shape_info.get("inputs", []) output_info = shape_info.get("outputs", []) + tensor_inputs = [value for value in inputs if isinstance(value, torch.Tensor)] + if not tensor_inputs: + raise RuntimeError( + "[torch.ops.tensorrt.execute_engine]: At least one tensor input is required to infer the engine output device" + ) + output_device = tensor_inputs[0].device + fake_mode = detect_fake_mode(inputs) if fake_mode is None: # No fake mode - shouldn't happen, but fall back to concrete shapes @@ -45,7 +52,7 @@ def _apply_symbolic_shape_expressions( for s in info["shape_exprs"] ] outputs.append( - torch.empty(shape, dtype=info["dtype"], device=inputs[0].device) + torch.empty(shape, dtype=info["dtype"], device=output_device) ) return outputs @@ -82,6 +89,8 @@ def in_compile_namespace(expr: sympy.Expr) -> sympy.Expr: # Align inputs: for each captured input, match it with the corresponding runtime input for inp_tensor, inp_info in zip(inputs, input_info): + if inp_info.get("is_scalar"): + continue for d, compile_expr in zip(inp_tensor.shape, inp_info["shape_exprs"]): if isinstance(compile_expr, int): continue @@ -226,7 +235,7 @@ def in_compile_namespace(expr: sympy.Expr) -> sympy.Expr: ) from e outputs.append( - torch.empty(output_shape, dtype=info["dtype"], device=inputs[0].device) + torch.empty(output_shape, dtype=info["dtype"], device=output_device) ) logger.debug( f"[torch.ops.tensorrt.execute_engine]: Meta kernel found the following output FakeTensors: {outputs}" diff --git a/tests/py/dynamo/models/test_exporter_inlining.py b/tests/py/dynamo/models/test_exporter_inlining.py index 2fa2f0470e..2029662569 100644 --- a/tests/py/dynamo/models/test_exporter_inlining.py +++ b/tests/py/dynamo/models/test_exporter_inlining.py @@ -6,6 +6,8 @@ import pytest import torch +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.symbolic_shapes import ShapeEnv from torch.export.graph_signature import ( ExportGraphSignature, InputKind, @@ -17,6 +19,7 @@ from torch_tensorrt.dynamo._exporter import ( create_trt_exp_program, inline_torch_modules, + inline_trt_modules, lift, ) @@ -147,6 +150,96 @@ def test_inline_torch_modules_computed_intermediate_inputs(): assert torch.allclose(out, torch.tensor(16.0)) +class _FakeTRTModule(torch.nn.Module): + def __init__(self, shape_info): + super().__init__() + self.engine = object() + self.symbolic_shape_expressions = shape_info + self.target_device = torch.device("cuda") + self.aliased_io = {} + + def forward(self, *args): + raise AssertionError("synthetic TRT module should be inlined, not executed") + + +@pytest.mark.unit +def test_inline_trt_modules_tensorizes_scalar_inputs_and_restores_outputs(): + """The inlined execute_engine Tensor[] boundary contains tensors only. + + Scalar partition inputs are materialized as zero-dimensional tensors before + execute_engine, and outputs recorded as scalars are converted back with item + before reaching their original consumers. + """ + shape_info = { + "inputs": [ + { + "shape_exprs": [], + "dtype": torch.int64, + "name": "length", + "is_scalar": True, + }, + { + "shape_exprs": [4], + "dtype": torch.float32, + "name": "x", + }, + ], + "outputs": [ + { + "shape_exprs": [], + "dtype": torch.int64, + "is_scalar": True, + }, + { + "shape_exprs": [4], + "dtype": torch.float32, + }, + ], + } + + shape_env = ShapeEnv() + fake_mode = FakeTensorMode(shape_env=shape_env) + with fake_mode: + scalar_val = shape_env.create_unbacked_symint() + tensor_val = torch.empty(4, device="cuda") + + graph = torch.fx.Graph() + scalar = graph.placeholder("length") + scalar.meta["val"] = scalar_val + tensor = graph.placeholder("x") + tensor.meta["val"] = tensor_val + + root = torch.nn.Module() + root.add_module("_run_on_acc_0", _FakeTRTModule(shape_info)) + call = graph.call_module("_run_on_acc_0", (scalar, tensor)) + call.meta["val"] = [scalar_val, tensor_val] + scalar_output = graph.call_function(operator.getitem, (call, 0)) + scalar_output.meta["val"] = scalar_val + tensor_output = graph.call_function(operator.getitem, (call, 1)) + tensor_output.meta["val"] = tensor_val + graph.output((scalar_output, tensor_output)) + gm = torch.fx.GraphModule(root, graph) + + inline_trt_modules(gm, expose_aliased_mutations=False) + gm.graph.lint() + + execute = next( + node + for node in gm.graph.nodes + if node.target is torch.ops.tensorrt.execute_engine.default + ) + engine_inputs = execute.args[0] + assert engine_inputs[0].target is torch.ops.aten.scalar_tensor.default + assert all(isinstance(node.meta.get("val"), torch.Tensor) for node in engine_inputs) + assert all(isinstance(value, torch.Tensor) for value in execute.meta["val"]) + + item_nodes = [ + node for node in gm.graph.nodes if node.target is torch.ops.aten.item.default + ] + assert len(item_nodes) == 1 + assert item_nodes[0].meta["val"] is scalar_val + + @pytest.mark.unit def test_create_trt_exp_program_rebuilds_in_spec_without_inputs(): """create_trt_exp_program must rebuild a correct in_spec on the plain-CodeGen diff --git a/tests/py/dynamo/models/test_meta_kernel_shape_inference.py b/tests/py/dynamo/models/test_meta_kernel_shape_inference.py index a0d6f731fd..4230c9b0ac 100644 --- a/tests/py/dynamo/models/test_meta_kernel_shape_inference.py +++ b/tests/py/dynamo/models/test_meta_kernel_shape_inference.py @@ -484,6 +484,58 @@ def test_rejects_inconsistent_repeated_direct_mapping(self): with pytest.raises(RuntimeError): _apply_symbolic_shape_expressions([fake_x, fake_z], shape_info) + @staticmethod + def _scalar_first_shape_info(): + return { + "inputs": [ + { + "shape_exprs": [], + "dtype": torch.int64, + "name": "length", + "is_scalar": True, + }, + { + "shape_exprs": [4], + "dtype": torch.float32, + "name": "x", + }, + ], + "outputs": [ + { + "shape_exprs": [4], + "dtype": torch.float32, + } + ], + } + + def test_scalar_input_is_not_walked_as_a_tensor(self): + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + scalar = shape_env.create_unbacked_symint() + tensor = torch.empty(4, device="cuda") + output = _apply_symbolic_shape_expressions( + [scalar, tensor], self._scalar_first_shape_info() + )[0] + + assert output.shape == (4,) + + def test_output_device_comes_from_first_tensor_input(self): + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + scalar = shape_env.create_unbacked_symint() + tensor = torch.empty(4, device="cuda") + fake_output = _apply_symbolic_shape_expressions( + [scalar, tensor], self._scalar_first_shape_info() + )[0] + + concrete_output = _apply_symbolic_shape_expressions( + [shape_env.create_unbacked_symint(), torch.empty(4)], + self._scalar_first_shape_info(), + )[0] + + assert fake_output.device == tensor.device + assert concrete_output.device.type == "cpu" + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) From 8bd1d57e151408d988a60b12af0a1e4fa311537f Mon Sep 17 00:00:00 2001 From: apbose Date: Tue, 1 Sep 2026 19:47:47 -0700 Subject: [PATCH 2/2] fix(dynamo): materialize SymFloat scalar inputs with the engine's binding dtype and rank torch.SymFloat inputs were recorded as float64, disagreeing with the float32 binding construct_submodule_inputs actually builds the engine with, and materialized as rank-0 tensors against its rank-1 binding -- so the exported graph disagreed with the real engine interface. Also guard against Input.dtype.unknown so ordinary tensor inputs without an explicit dtype don't crash the new dtype lookup. --- .../dynamo/conversion/_conversion.py | 7 +- .../conversion/_symbolic_shape_capture.py | 30 ++++++-- .../dynamo/partitioning/common.py | 4 +- .../conversion/test_symbolic_shape_capture.py | 76 +++++++++++++++++++ .../dynamo/models/test_symint_scalar_input.py | 53 +++++++++++++ 5 files changed, 161 insertions(+), 9 deletions(-) create mode 100644 tests/py/dynamo/conversion/test_symbolic_shape_capture.py diff --git a/py/torch_tensorrt/dynamo/conversion/_conversion.py b/py/torch_tensorrt/dynamo/conversion/_conversion.py index 38d104473a..f70b732cbd 100644 --- a/py/torch_tensorrt/dynamo/conversion/_conversion.py +++ b/py/torch_tensorrt/dynamo/conversion/_conversion.py @@ -4,6 +4,7 @@ import logging from typing import Any, Dict, List, NamedTuple, Optional, Sequence, Tuple +import tensorrt as trt import torch from torch_tensorrt._enums import dtype from torch_tensorrt._features import ENABLED_FEATURES @@ -25,8 +26,6 @@ ) from torch_tensorrt.logging import TRT_LOGGER -import tensorrt as trt - logger = logging.getLogger(__name__) @@ -221,7 +220,9 @@ def interpret_module_to_result( SerializedInterpreterResult """ - symbolic_shape_expressions = extract_symbolic_shape_expressions(module) + symbolic_shape_expressions = extract_symbolic_shape_expressions( + module, inputs=inputs + ) if symbolic_shape_expressions is None: raise RuntimeError( "Failed to extract symbolic shape expressions from source FX graph partition" diff --git a/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py b/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py index e8f84921a2..11213d9d98 100644 --- a/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py +++ b/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py @@ -7,15 +7,18 @@ """ import logging -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Sequence import torch +from torch_tensorrt._enums import dtype as _dtype +from torch_tensorrt._Input import Input logger = logging.getLogger(__name__) def extract_symbolic_shape_expressions( module: torch.fx.GraphModule, + inputs: Optional[Sequence[Input]] = None, ) -> Optional[Dict[str, List[Dict[str, Any]]]]: """ Extract symbolic shape expressions from an FX graph. @@ -25,11 +28,23 @@ def extract_symbolic_shape_expressions( Args: module: FX GraphModule with symbolic shapes in node metadata + inputs: Engine Input specs from construct_submodule_inputs, aligned by + name. Used as the dtype source of truth for scalar inputs, which + have no dtype of their own in FX metadata. Falls back to a + best-effort default when not provided. Returns: Dict with 'inputs' and 'outputs' keys, each containing a list of dicts with shape_exprs and dtype, or None if extraction fails """ + # dtype.unknown has no torch.dtype equivalent; skip it rather than raise + # (unset-dtype inputs are never looked up below anyway). + input_dtypes_by_name = { + inp.name: inp.dtype.to(torch.dtype) + for inp in inputs or () + if inp.dtype != _dtype.unknown + } + # Find input nodes (placeholders) input_nodes = [node for node in module.graph.nodes if node.op == "placeholder"] @@ -70,11 +85,15 @@ def extract_symbolic_shape_expressions( ) elif isinstance(input_val, (torch.SymInt, torch.SymFloat, int, float, bool)): if isinstance(input_val, (torch.SymInt, int)): - scalar_dtype = torch.int64 + default_scalar_dtype = torch.int64 elif isinstance(input_val, (torch.SymFloat, float)): - scalar_dtype = torch.float64 + default_scalar_dtype = torch.float32 else: - scalar_dtype = torch.bool + default_scalar_dtype = torch.bool + # Prefer the engine's actual binding dtype over the guess above. + scalar_dtype = input_dtypes_by_name.get( + input_node.name, default_scalar_dtype + ) input_info.append( { "shape_exprs": [], @@ -122,7 +141,8 @@ def extract_symbolic_shape_expressions( if isinstance(out_val, (torch.SymInt, int)): scalar_dtype = torch.int64 elif isinstance(out_val, (torch.SymFloat, float)): - scalar_dtype = torch.float64 + # No float64 output binding exists in TensorRT. + scalar_dtype = torch.float32 else: scalar_dtype = torch.bool output_info.append( diff --git a/py/torch_tensorrt/dynamo/partitioning/common.py b/py/torch_tensorrt/dynamo/partitioning/common.py index 62e68449b6..a0942a5d0b 100644 --- a/py/torch_tensorrt/dynamo/partitioning/common.py +++ b/py/torch_tensorrt/dynamo/partitioning/common.py @@ -283,9 +283,11 @@ def construct_submodule_inputs( ) ) elif isinstance(input_meta, torch.SymFloat): + # Rank-0 to match the 0-D scalar_tensor inline_trt_modules + # materializes for this input at the engine boundary. torchtrt_inputs.append( get_input( - [1], + [], torch.float32, name=input.name, is_shape_tensor=False, # Only SymInt inputs are treated as shape tensors diff --git a/tests/py/dynamo/conversion/test_symbolic_shape_capture.py b/tests/py/dynamo/conversion/test_symbolic_shape_capture.py new file mode 100644 index 0000000000..8a84661117 --- /dev/null +++ b/tests/py/dynamo/conversion/test_symbolic_shape_capture.py @@ -0,0 +1,76 @@ +import torch +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.symbolic_shapes import ShapeEnv +from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt import Input +from torch_tensorrt.dynamo.conversion._symbolic_shape_capture import ( + extract_symbolic_shape_expressions, +) + + +class TestSymbolicShapeCaptureScalarDtype(TestCase): + """A SymFloat scalar input has no dtype of its own in FX metadata; the + recorded dtype must match the engine's actual binding (float32).""" + + def _make_symfloat_input_module(self) -> torch.fx.GraphModule: + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + scale = shape_env.create_unbacked_symfloat() + x = torch.empty((4,), dtype=torch.float32) + + graph = torch.fx.Graph() + scale_input = graph.placeholder("scale") + scale_input.meta["val"] = scale + x_input = graph.placeholder("x") + x_input.meta["val"] = x + output = graph.call_function( + torch.ops.aten.mul.Tensor, args=(x_input, scale_input) + ) + output.meta["val"] = x + graph.output(output) + return torch.fx.GraphModule({}, graph) + + def test_symfloat_input_uses_engine_binding_dtype(self): + module = self._make_symfloat_input_module() + engine_inputs = [ + Input([1], dtype=torch.float32, name="scale"), + Input((4,), dtype=torch.float32, name="x"), + ] + + metadata = extract_symbolic_shape_expressions(module, inputs=engine_inputs) + + scale_info = next( + info for info in metadata["inputs"] if info["name"] == "scale" + ) + self.assertTrue(scale_info["is_scalar"]) + self.assertEqual(scale_info["dtype"], torch.float32) + + def test_symfloat_input_defaults_to_float32_without_engine_inputs(self): + # Used to hardcode float64 -- no real engine binding is ever float64. + module = self._make_symfloat_input_module() + + metadata = extract_symbolic_shape_expressions(module) + + scale_info = next( + info for info in metadata["inputs"] if info["name"] == "scale" + ) + self.assertEqual(scale_info["dtype"], torch.float32) + + def test_ordinary_input_without_explicit_dtype_does_not_raise(self): + # An Input with dtype omitted (dtype.unknown) used to raise TypeError. + module = self._make_symfloat_input_module() + engine_inputs = [ + Input([1], dtype=torch.float32, name="scale"), + Input((4,)), # dtype intentionally omitted + ] + + metadata = extract_symbolic_shape_expressions(module, inputs=engine_inputs) + + scale_info = next( + info for info in metadata["inputs"] if info["name"] == "scale" + ) + self.assertEqual(scale_info["dtype"], torch.float32) + + +if __name__ == "__main__": + run_tests() diff --git a/tests/py/dynamo/models/test_symint_scalar_input.py b/tests/py/dynamo/models/test_symint_scalar_input.py index 9ab3104064..66f366c05c 100644 --- a/tests/py/dynamo/models/test_symint_scalar_input.py +++ b/tests/py/dynamo/models/test_symint_scalar_input.py @@ -14,6 +14,7 @@ import pytest import torch import torch_tensorrt as torchtrt +from torch_tensorrt.dynamo._exporter import transform from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity assertions = unittest.TestCase() @@ -194,3 +195,55 @@ def forward(self, x, targets): ) torch._dynamo.reset() + + +@pytest.mark.unit +def test_symfloat_scalar_input(): + """ + A data-dependent Python float (e.g. weights.sum().item()) crossing a TRT + partition boundary as a SymFloat must be materialized with the engine's + actual binding dtype (float32), not a guessed one. + """ + torch._dynamo.config.capture_scalar_outputs = True + try: + + class ScaleByDataDependentFloat(torch.nn.Module): + def forward(self, x, weights): + scale = weights.sum().item() + return x * scale + + model = ScaleByDataDependentFloat().eval().cuda() + x = torch.randn(8).cuda() + weights = torch.randn(4).cuda() + expected = model(x, weights) + + exported = torch.export.export(model, (x, weights)) + compiled = torchtrt.dynamo.compile( + exported, + inputs=(x, weights), + min_block_size=1, + pass_through_build_failures=True, + ) + + # The inserted scalar tensor must carry the engine binding's dtype. + inlined = transform(compiled) + inlined.recompile() + scalar_tensor_nodes = [ + node + for node in inlined.graph.nodes + if node.target == torch.ops.aten.scalar_tensor.default + ] + assertions.assertEqual(len(scalar_tensor_nodes), 1) + assertions.assertEqual( + scalar_tensor_nodes[0].kwargs["dtype"], + torch.float32, + msg="scalar_tensor inserted for the SymFloat input must match the " + "float32 engine binding dtype, not the Python-type-based default", + ) + + # And the exported module must actually execute successfully. + actual = compiled(x, weights) + torch.testing.assert_close(actual, expected) + finally: + torch._dynamo.config.capture_scalar_outputs = False + torch._dynamo.reset()