fix(dynamo): several correctness and speed fixes found compiling large models

py/torch_tensorrt/dynamo/_compiler.py

@@ -225,8 +225,9 @@ def cross_compile_for_windows(
         disabled_constant_fold_exclusions (Collection[str]): IDs of predefined
             Torch-TensorRT rules to disable. Rules are registered by lowering
             implementations and enabled by default; compilation users do not
-            need to register them. Disabling a rule makes its matching nodes
-            foldable. Supported IDs: ``"attention_mask_arange"``.
+            need to register them. Disabling a rule makes the nodes that only
+            it matches foldable. Supported IDs: ``"attention_mask_arange"``,
+            ``"scatter_add_positions"``.
             Default is empty.
         attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False.
         fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False.
@@ -428,12 +429,16 @@ def cross_compile_for_windows(
 
     # Apply lowering on the graph module. Note: constant_fold runs inside post_lowering and requires
     # module parameters to still be on GPU, so we must not deallocate before this call.
+    unlowered_gm = gm
     gm = post_lowering(gm, settings)
     logger.debug(f"CPU memory usage after post_lowering: {get_cpu_memory_usage()} MB")
     logger.debug("Lowered Input graph: " + str(gm.graph))
 
     # Move the weights in the state_dict to CPU
     if offload_module_to_cpu:
+        # Complex lowering can return a new module without the weights that only
+        # constant folding read, so offload the module lowering started from too.
+        deallocate_module(unlowered_gm)
         deallocate_module(gm)
         logger.info(
             "The PyTorch model was moved to the CPU to allocate all GPU memory to TensorRT. To retain the model on the GPU, set offload_module_to_cpu=False"
@@ -637,8 +642,9 @@ def compile(
         disabled_constant_fold_exclusions (Collection[str]): IDs of predefined
             Torch-TensorRT rules to disable. Rules are registered by lowering
             implementations and enabled by default; compilation users do not
-            need to register them. Disabling a rule makes its matching nodes
-            foldable. Supported IDs: ``"attention_mask_arange"``.
+            need to register them. Disabling a rule makes the nodes that only
+            it matches foldable. Supported IDs: ``"attention_mask_arange"``,
+            ``"scatter_add_positions"``.
             Default is empty.
         attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False.
         fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False.
@@ -909,6 +915,7 @@ def compile(
 
     # Apply lowering on the graph module. Note: constant_fold runs inside post_lowering and requires
     # module parameters to still be on GPU, so we must not deallocate before this call.
+    unlowered_gm = gm
     gm = post_lowering(gm, settings)
     logger.debug(f"CPU memory usage after post_lowering: {get_cpu_memory_usage()} MB")
     logger.debug("Lowered Input graph: " + str(gm.graph))
@@ -921,6 +928,9 @@ def compile(
 
     # Move the weights in the state_dict to CPU
     if offload_module_to_cpu:
+        # Complex lowering can return a new module without the weights that only
+        # constant folding read, so offload the module lowering started from too.
+        deallocate_module(unlowered_gm)
         deallocate_module(gm)
         logger.info(
             "The PyTorch model was moved to the CPU to allocate all GPU memory to TensorRT. To retain the model on the GPU, set offload_module_to_cpu=False"
@@ -2102,8 +2112,9 @@ def convert_exported_program_to_serialized_trt_engine(
         disabled_constant_fold_exclusions (Collection[str]): IDs of predefined
             Torch-TensorRT rules to disable. Rules are registered by lowering
             implementations and enabled by default; compilation users do not
-            need to register them. Disabling a rule makes its matching nodes
-            foldable. Supported IDs: ``"attention_mask_arange"``.
+            need to register them. Disabling a rule makes the nodes that only
+            it matches foldable. Supported IDs: ``"attention_mask_arange"``,
+            ``"scatter_add_positions"``.
             Default is empty.
         attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False.
         fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False.

py/torch_tensorrt/dynamo/_refit.py

@@ -460,9 +460,7 @@ def refit_module_weights(
 
     if verify_output and arg_inputs is not None:
         new_gm.to(to_torch_device(settings.device))  # move to device for inference
-        # A complex placeholder stays complex in the lowered reference module; the
-        # unpacking to real happens inside the graph, so both modules take the same
-        # inputs the caller passed.
+        _restore_complex_io(new_gm)
         outputs_match = check_module_output(
             new_module=new_gm,
             refitted_module=compiled_module,
@@ -479,6 +477,42 @@ def refit_module_weights(
     return compiled_module
 
 
+def _restore_complex_io(gm: torch.fx.GraphModule) -> None:
+    """Make the lowered reference module take and return complex tensors, as the
+    compiled module does, so the output check can compare the two.
+
+    The legacy complex rewriter replaces a complex input with a real
+    "<name>_unpacked_complex" one, and both complex lowerings return a complex
+    output as a real tensor with a trailing dimension of 2.
+    """
+    graph = gm.graph
+    unpacked = {f"{n}_unpacked_complex" for n in gm.meta.get("complex_input_names", [])}
+    for node in list(graph.nodes):
+        if node.op == "placeholder" and node.name in unpacked:
+            with graph.inserting_after(node):
+                real = graph.call_function(torch.ops.aten.view_as_real.default, (node,))
+            node.replace_all_uses_with(real)
+            real.args = (node,)
+    output_node = list(graph.nodes)[-1]
+    output_arg = output_node.args[0]
+    outputs = (
+        list(output_arg) if isinstance(output_arg, (list, tuple)) else [output_arg]
+    )
+    for idx in gm.meta.get("complex_output_indices", []):
+        with graph.inserting_before(output_node):
+            outputs[idx] = graph.call_function(
+                torch.ops.aten.view_as_complex.default, (outputs[idx],)
+            )
+    output_node.args = (
+        (
+            type(output_arg)(outputs)
+            if isinstance(output_arg, (list, tuple))
+            else outputs[0]
+        ),
+    )
+    gm.recompile()
+
+
 # Util functions -----------
 import base64
 

py/torch_tensorrt/dynamo/_settings.py

@@ -159,8 +159,9 @@ class CompilationSettings:
         disabled_constant_fold_exclusions (Collection[str]): IDs of predefined
             Torch-TensorRT rules to disable. Rules are registered by lowering
             implementations and enabled by default; compilation users do not
-            need to register them. Disabling a rule makes its matching nodes
-            foldable. Supported IDs: ``"attention_mask_arange"``.
+            need to register them. Disabling a rule makes the nodes that only
+            it matches foldable. Supported IDs: ``"attention_mask_arange"``,
+            ``"scatter_add_positions"``.
             Default is empty.
         attn_bias_is_causal (bool): Whether the attn_bias in efficient SDPA is causal. Default is True. This can accelerate models from HF because attn_bias is always a causal mask in HF. If you want to use non-causal attn_bias, you can set this to False.
         fold_shared_kv_attention (bool): For attention where all query heads share one key/value head (multi-query attention), compute it against that one head instead of the copies that Hugging Face models make for each query head. TensorRT then runs two plain matrix products instead of its fused attention kernel, which can be faster or slower depending on the GPU, so measure it. In float16 the result can round differently; use_fp32_acc=True reduces that. Only static shapes are rewritten, and scaled_dot_product_attention also needs decompose_attention=True. Both preserved matmuls and the reshape/bmm form produced by AOT tracing under torch.compile are supported for repeat_kv style graphs. In the bmm form, broadcasting a single KV head with enable_gqa=True is not folded. Default is False.

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -1753,7 +1753,27 @@ def aten_ops_scatter(
     )
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.select.int, supports_dynamic_shapes=True)
+def select_validator(
+    node: Node, settings: Optional[CompilationSettings] = None
+) -> bool:
+    """Keep a select of a uint8 tensor on the PyTorch path.
+
+    TensorRT allows uint8 only at engine inputs and at outputs made by a cast or
+    identity layer, so the gather's uint8 result fails the engine build, as it does
+    for aten.gather.
+    """
+    meta = getattr(args_bounds_check(node.args, 0), "meta", {})
+    if getattr(meta.get("val", meta.get("tensor_meta")), "dtype", None) == torch.uint8:
+        _LOGGER.debug("select of a uint8 tensor is not supported, falling back")
+        return False
+    return True
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten.select.int,
+    capability_validator=select_validator,
+    supports_dynamic_shapes=True,
+)
 def aten_ops_select(
     ctx: ConversionContext,
     target: Target,
@@ -2142,12 +2162,45 @@ def aten_ops_permute(
     )
 
 
+def _truncatable_float64_copy(node: Node) -> bool:
+    """Whether a float64 copy can run as a float32 copy and give PyTorch's result.
+
+    The source has to be a float type TensorRT casts to float32, and the PyTorch
+    result has to have the contiguous layout a TensorRT copy produces. The check
+    does not depend on where the copy sits, because partitioning can turn any copy
+    into an engine output and the converter is validated again then.
+    """
+    input_node = node.args[0] if node.args else None
+    input_meta = input_node.meta.get("val") if isinstance(input_node, Node) else None
+    if not isinstance(input_meta, torch.Tensor) or input_meta.dtype not in (
+        torch.float16,
+        torch.bfloat16,
+        torch.float32,
+        torch.float64,
+    ):
+        return False
+    # Alone on a placeholder after partitioning, a dynamic copy must also pass
+    # _dynamic_placeholder_copy_supported, which refuses any explicit memory_format.
+    if node.kwargs.get("memory_format") is not None and node_has_dynamic_shapes(node):
+        return False
+    output_meta = node.meta.get("val")
+    if not isinstance(output_meta, torch.Tensor):
+        return False
+    # A data-dependent stride cannot be decided here, so leave the copy to PyTorch.
+    try:
+        return bool(output_meta.is_contiguous())
+    except Exception:
+        return False
+
+
 def to_copy_dtype_validator(
     placeholder_only: bool, settings: Optional[CompilationSettings] = None
 ) -> Callable[[Node, CompilationSettings], bool]:
     """Return validator for to_copy node with placeholder restrictions"""
 
-    def validate_dtype(to_copy_node: Node) -> bool:
+    def validate_dtype(
+        to_copy_node: Node, settings: Optional[CompilationSettings]
+    ) -> bool:
         """Returns true if the to_copy node can be converted to TRT
 
         Based on data type being casted to
@@ -2187,6 +2240,15 @@ def validate_dtype(to_copy_node: Node) -> bool:
         if "dtype" in to_copy_node.kwargs:
             if to_copy_node.kwargs["dtype"] in allowed_casts:
                 return True
+            # truncate_double means the caller accepts float32 for float64, and
+            # the converter maps it.
+            elif (
+                to_copy_node.kwargs["dtype"] is torch.float64
+                and settings is not None
+                and settings.truncate_double
+                and _truncatable_float64_copy(to_copy_node)
+            ):
+                return True
             else:
                 _LOGGER.debug(
                     f"_to_copy converter rejected node {to_copy_node} with dtype {to_copy_node.kwargs['dtype']}"
@@ -2206,7 +2268,7 @@ def validator(
         """
         # The placeholder restriction is satsfied if placeholder_only is the same
         # truth value as is_only_operator_on_placeholder(to_copy_node)
-        return validate_dtype(to_copy_node) and (
+        return validate_dtype(to_copy_node, settings) and (
             (not placeholder_only) ^ is_only_operator_on_placeholder(to_copy_node)
         )
 

py/torch_tensorrt/dynamo/conversion/impl/cast.py

@@ -34,6 +34,16 @@ def to_copy(
     if not isinstance(input, TRTTensor):
         input = get_trt_tensor(ctx, input, f"{name}_copy_tensor")
 
+    # TensorRT has no float64. With truncate_double the caller accepts float32,
+    # which holds these source types exactly but not every larger integer.
+    if (
+        ctx.compilation_settings.truncate_double
+        and _enums.dtype._from(dtype) == _enums.dtype.f64
+        and input.dtype
+        in {trt.bool, trt.int8, trt.uint8, trt.float16, trt.bfloat16, trt.float32}
+    ):
+        dtype = torch.float32
+
     # If cast is forced, insert identity layer regardless of whether the dtype
     # doesn't change
     if force_layer:

py/torch_tensorrt/dynamo/conversion/impl/select.py

@@ -683,6 +683,29 @@ def index_put_scatter_add_plugin(
     return out
 
 
+def _broadcast_target(
+    expected_shape: Sequence[Union[int, ITensor]], tensor_shape: Sequence[int]
+) -> List[Union[int, ITensor]]:
+    """The shape to expand a tensor laid out as (N, ...) to, broadcasting like PyTorch.
+
+    expand broadcasts any static dim whose target is an ITensor, which is only right
+    for size 1. A data-dependent N, such as the True count of a boolean mask, must
+    match a larger static first dim at runtime in PyTorch, so keep that dim.
+    """
+    tensor_shape = [1] * (len(expected_shape) - len(tensor_shape)) + list(tensor_shape)
+    shape: List[Union[int, ITensor]] = []
+    for k, (exp_dim, dim) in enumerate(zip(expected_shape, tensor_shape)):
+        if k == 0 and isinstance(exp_dim, ITensor) and dim not in (1, DYNAMIC_DIM):
+            shape.append(dim)
+        elif dim in (1, DYNAMIC_DIM) or exp_dim == dim:
+            shape.append(exp_dim)
+        else:
+            raise ValueError(
+                f"Cannot broadcast {tuple(tensor_shape)} to {expected_shape}"
+            )
+    return shape
+
+
 def index_put_converter(
     ctx: ConversionContext,
     target: Target,
@@ -853,7 +876,7 @@ def index_put_converter(
             source_ir,
             f"{name}_expand_idx_I_{i}",
             idx_reshaped,
-            (N, F_volume),
+            tuple(_broadcast_target((N, F_volume), idx_reshaped.shape)),
         )
         I_tensors.append(expanded_idx)
 
@@ -999,8 +1022,7 @@ def index_put_converter(
 
     # Define the expected shape based on (N,) + F_shape_values.
     # N may be a ITensor when it comes from a dynamic source (e.g. nonzero).
-    # Pass it directly — impl.slice.expand handles ITensor shape elements and
-    # will emit a stride-0 broadcast when expanding a size-1 dim to a dynamic N.
+    # _broadcast_target decides which values dims expand to it.
     expected_shape = (N,) + tuple(F_shape_values)
 
     # Broadcast 'values' to match the expected shape
@@ -1037,7 +1059,7 @@ def index_put_converter(
                 source_ir,
                 f"{name}_expand_values",
                 values_permuted,
-                expected_shape,
+                tuple(_broadcast_target(expected_shape, values_permuted.shape)),
             )
         elif K == 1 and len(values.shape) == 1 and F and I[0] > max(F):
             # For data[..., idx] = values, a 1-D values tensor carries the
@@ -1114,27 +1136,13 @@ def index_put_converter(
             )
         else:
             # Discontinuous case (K > 1 or K == 0)
-            values_shape_padded = [1] * (
-                len(expected_shape) - len(values.shape)
-            ) + list(values.shape)
-            broadcast_shape = []
-            for exp_dim, val_dim in zip(expected_shape, values_shape_padded):
-                if val_dim == DYNAMIC_DIM or exp_dim == DYNAMIC_DIM:
-                    broadcast_shape.append(-1)
-                elif val_dim == 1 or exp_dim == val_dim:
-                    broadcast_shape.append(exp_dim)
-                else:
-                    raise ValueError(
-                        f"Cannot broadcast {values.shape} to {expected_shape}"
-                    )
-
             values_expanded = impl.slice.expand(
                 ctx,
                 target,
                 source_ir,
                 f"{name}_expand_values",
                 values,
-                expected_shape,
+                tuple(_broadcast_target(expected_shape, values.shape)),
             )
 
     # Flatten values to (N * F_volume,)