fix(dynamo): normalize scalar values at engine boundaries

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

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"

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(

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

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]
 

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}"