Skip installing large aliased constant-folded weights

py/torch_tensorrt/dynamo/_compiler.py

@@ -5,7 +5,6 @@
 import os
 import platform
 import warnings
-
 from typing import Any, Collection, Dict, List, Optional, Sequence, Tuple, Union
 
 import sympy
@@ -1263,6 +1262,9 @@ def preserve_module_specs(
                 torch_executed_ops=settings.torch_executed_ops,
                 require_full_compilation=settings.require_full_compilation,
                 skip_fusion=(num_supported_ops == total_ops),
+                assume_full_support=(
+                    settings.require_full_compilation and num_supported_ops == total_ops
+                ),
             )
 
         except torch.fx.passes.splitter_base.FxNetSplitterInternalError:

py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py

@@ -18,6 +18,30 @@
 
 logger = logging.getLogger(__name__)
 
+# Skip installing large folded tensors that already share storage with an
+# existing module attribute (typically weight permutes/views). On Flux NVFP4
+# those dominate replace cost via cpu().contiguous() under offload_module_to_cpu.
+# Do NOT skip materialized folds (e.g. VLA position embedding outputs) — those
+# can be required for TRT legality (int64 indices never reach the converter).
+_MAX_CONSTANT_FOLD_BYTES = 1 << 20  # 1 MiB
+
+
+def _tensor_reuses_module_storage(
+    gm: torch.fx.GraphModule, constant: torch.Tensor
+) -> bool:
+    try:
+        ptr = constant.untyped_storage().data_ptr()
+    except Exception:
+        return False
+    for tensor in gm.state_dict().values():
+        if isinstance(tensor, torch.Tensor):
+            try:
+                if tensor.untyped_storage().data_ptr() == ptr:
+                    return True
+            except Exception:
+                continue
+    return False
+
 
 @torch.utils._python_dispatch._disable_current_modes()  # type: ignore
 def constant_fold(
@@ -35,17 +59,38 @@ def constant_fold(
 
     # The constants are created on CPU to save GPU memory for TensorRT compilation.
     # For TRT INetwork construction the constants are moved to CPU in get_attr call.
+    skipped_alias = 0
     for node, constant in cf.node_replacements.items():
+        if isinstance(constant, torch.Tensor):
+            nbytes = int(constant.numel() * constant.element_size())
+            if nbytes > _MAX_CONSTANT_FOLD_BYTES and _tensor_reuses_module_storage(
+                gm, constant
+            ):
+                skipped_alias += 1
+                logger.debug(
+                    "Skipping constant-fold install for aliased %s (%d bytes > %d)",
+                    node.name,
+                    nbytes,
+                    _MAX_CONSTANT_FOLD_BYTES,
+                )
+                continue
+        # Register folded values as plain tensors (buffers), matching Inductor.
         if settings.offload_module_to_cpu:
             replace_node_with_constant(
                 gm,
                 node,
-                torch.nn.Parameter(constant.cpu().contiguous(), requires_grad=False),
+                constant.cpu().contiguous(),
             )
         else:
-            replace_node_with_constant(
-                gm, node, torch.nn.Parameter(constant, requires_grad=False)
-            )
+            replace_node_with_constant(gm, node, constant)
+
+    if skipped_alias:
+        logger.info(
+            "Skipped installing %d large aliased folded constant(s) (>%d bytes); "
+            "leaving original view/permute ops in the graph",
+            skipped_alias,
+            _MAX_CONSTANT_FOLD_BYTES,
+        )
 
     erased_params = []
     for node in gm.graph.nodes:
@@ -71,7 +116,7 @@ def replace_node_with_constant(
     """Adapted from:
     https://github.com/pytorch/pytorch/blob/bcf35c6ae62bb6560befa3550e37a8283944e5f4/torch/_inductor/constant_folding.py#L17-L43
 
-    Modified to register parameters, instead of buffers for frozen constants
+    Registers frozen constants as buffers (same as Inductor), not Parameters.
     """
     g = gm.graph
 
@@ -95,7 +140,7 @@ def replace_node_with_constant(
         g.erase_node(node)
 
     # Needed to suppress `does not reference an nn.Module, nn.Parameter, or buffer` warning
-    gm.register_parameter(qualname, constant)
+    gm.register_buffer(qualname, constant)
     setattr(gm, qualname, constant)
 
 

py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py

@@ -138,6 +138,7 @@ def __init__(
         require_full_compilation: bool = REQUIRE_FULL_COMPILATION,
         return_tuple: bool = False,
         skip_fusion: bool = False,
+        assume_full_support: bool = False,
     ):
         """
         Preprocesses graph before splitting:
@@ -157,20 +158,30 @@ def __init__(
             skip_fusion=skip_fusion,
         )
         self.operator_support = operator_support
-
-        # Get all accelerated nodes based on operator support conditions
-        self.acc_nodes = FxNetAccNodesFinder(
-            self.module, self.operator_support, self.settings.allow_non_tensor
-        )()
-
-        if self.settings.skip_fusion:
-            self.fusions = {}
+        self.assume_full_support = assume_full_support
+
+        if self.assume_full_support:
+            # The caller already walked the graph and verified converter support.
+            # Avoid repeating support discovery and dependency construction.
+            self.acc_nodes = {
+                node for node in self.module.graph.nodes if node.op in CALLABLE_NODE_OPS
+            }
+            self.fusions: Dict[torch.fx.Node, NodeSet] = {}
+            self.deps: Dict[torch.fx.Node, NodeSet] = {}
         else:
-            self.fusions = FxNetAccFusionsFinder(module, set(self.acc_nodes))()
+            # Get all accelerated nodes based on operator support conditions
+            self.acc_nodes = FxNetAccNodesFinder(
+                self.module, self.operator_support, self.settings.allow_non_tensor
+            )()
 
-        # Modify deps to add more deps for fused nodes
-        self.deps = self.find_deps()
-        self.update_deps_for_fusions()
+            if self.settings.skip_fusion:
+                self.fusions = {}
+            else:
+                self.fusions = FxNetAccFusionsFinder(module, set(self.acc_nodes))()
+
+            # Modify deps to add more deps for fused nodes
+            self.deps = self.find_deps()
+            self.update_deps_for_fusions()
 
         self.non_acc_submodule_name = "_run_on_gpu_"
         self._node_submodule_map: Dict[str, str] = {}
@@ -223,6 +234,40 @@ def partition_graph(self) -> torch.fx.GraphModule:
 
         Returns a GraphModule with submodules for each segment
         """
+        unsupported = getattr(self.operator_support, "unsupported_operators", None)
+        # The explicit assumption comes from the compiler's earlier support walk.
+        # Otherwise, an empty dict means AccNodesFinder found no unsupported ops.
+        fully_supported = self.assume_full_support or (
+            isinstance(unsupported, dict) and len(unsupported) == 0
+        )
+
+        # Fast path: user demanded a single TRT engine and every op is convertible.
+        # Skip adjacency splitting; emit one ACC block and tag/split as usual so the
+        # rest of compile still sees `_run_on_acc_*`.
+        if self.require_full_compilation and fully_supported:
+            if self.settings.min_acc_module_size != MIN_BLOCK_SIZE:
+                logger.warning(
+                    "Detected both require_full_compilation and min_block_size compilation "
+                    "arguments were specified. Disregarding min_block_size argument for "
+                    "fully supported model."
+                )
+            nodes = [
+                node for node in self.module.graph.nodes if node.op in CALLABLE_NODE_OPS
+            ]
+            if not nodes:
+                raise AssertionError(
+                    "require_full_compilation=True was specified, but no accelerated "
+                    "operators were found in the graph"
+                )
+            logger.info(
+                "require_full_compilation + full operator support: "
+                "wrapping graph as a single TRT submodule (skipping adjacency split)"
+            )
+            subgraphs = [Subgraph(is_acc=True, nodes=nodes)]
+            self.num_trt_accelerated_subgraphs = 1
+            self.tag(subgraphs)
+            return self.split(remove_tag=True)
+
         # Delegate nodes based on operator coverage
         subgraphs = self.put_nodes_into_subgraphs()
 
@@ -280,6 +325,7 @@ def partition(
     torch_executed_ops: Collection[Target] = set(),
     require_full_compilation: bool = REQUIRE_FULL_COMPILATION,
     skip_fusion: bool = False,
+    assume_full_support: bool = False,
 ) -> Tuple[torch.fx.GraphModule, OpSupportTester]:
     """Partition an FX GraphModule with aten ops into TRT engines
     Partitioning is based on converter operator support
@@ -290,6 +336,8 @@ def partition(
         torch_executed_ops: Collection of operations to run in Torch, regardless of converter coverage
         require_full_compilation: Require that all computational operators be run in TRT
         skip_fusion: Skip fusions found by FxNetAccFusionsFinder
+        assume_full_support: Skip repeated support/dependency discovery because
+            the caller already verified that every computational op is supported
     Returns:
         torch.fx.GraphModule, OpSupportTester
     """
@@ -306,6 +354,7 @@ def partition(
         min_block_size=min_block_size,
         require_full_compilation=require_full_compilation,
         skip_fusion=skip_fusion,
+        assume_full_support=assume_full_support,
     )
 
     partitioned_graph = partitioner.partition_graph()