Skip installing large aliased constant-folded weights

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)