pytorch
diff --git a/‎test/dynamo/test_misc.py
Lines changed: 48 additions & 0 deletions b/‎test/dynamo/test_misc.py
Lines changed: 48 additions & 0 deletions
diff --git a/‎torch/_dynamo/variables/builder.py
Lines changed: 16 additions & 2 deletions b/‎torch/_dynamo/variables/builder.py
Lines changed: 16 additions & 2 deletions
diff --git a/‎torch/compiler/config.py
Lines changed: 8 additions & 0 deletions b/‎torch/compiler/config.py
Lines changed: 8 additions & 0 deletions
@@ -7858,6 +7858,54 @@ def my_dyn_fn(x, y):
         with self.assertRaises(ConstraintViolationError):
             torch.compile(my_dyn_fn, backend="eager")(y, y)
 
+    @torch.compiler.config.patch(dynamic_sources="L['x']")
+    def test_dynamic_int_sources(self):
+        counter = CompileCounter()
+
+        @torch.compile(backend=counter)
+        def fn(x):
+            return torch.randn(5) * x
+
+        fn(1)
+        fn(2)
+        fn(3)
+
+        self.assertEqual(counter.frame_count, 1)
+
+    @torch.compiler.config.patch(dynamic_sources="L['x']")
+    def test_dynamic_tensor_sources(self):
+        counter = CompileCounter()
+
+        @torch.compile(backend=counter)
+        def fn(x):
+            return x * x
+
+        fn(torch.randn(2))
+        fn(torch.randn(3))
+        fn(torch.randn(4))
+
+        self.assertEqual(counter.frame_count, 1)
+
+    @torch.compiler.config.patch(dynamic_sources="L['x']")
+    def test_dynamic_sources_graph_break(self):
+        counter = CompileCounter()
+
+        def foo(x):
+            return x * x
+
+        @torch.compile(backend=counter)
+        def fn(x):
+            x = x * x
+            torch._dynamo.graph_break()
+            return foo(x)
+
+        fn(torch.randn(2))
+        fn(torch.randn(3))
+        fn(torch.randn(4))
+
+        # 2 since graph break produces 2 graphs. NB: there are no recompiles
+        self.assertEqual(counter.frame_count, 2)
+
     def test_cannot_trace_mark_dynamic(self):
         y = torch.randn([3, 3, 3])
 
 
@@ -275,6 +275,7 @@
 static_inputs_log = torch._logging.getArtifactLogger(
     __name__, "cudagraph_static_inputs"
 )
+_DYNAMIC_SOURCES: Optional[set[str]] = None
 
 
 DimList = list
@@ -1955,7 +1956,10 @@ def wrap_symint(self, value):
                 and frame_state_entry.scalar is auto_dynamic
             ):
                 dynamic_dim = get_automatic_dynamic_shapes_mark_as()
-            elif not config.assume_static_by_default:
+            elif (
+                not config.assume_static_by_default
+                or self.source.name() in get_dynamic_sources()
+            ):
                 dynamic_dim = DimDynamic.DYNAMIC
             else:  # assume_static_by_default
                 # TODO: dynamic_dim = DimDynamic.STATIC should work but
@@ -2639,6 +2643,15 @@ def get_automatic_dynamic_shapes_mark_as():
         )
 
 
+def get_dynamic_sources() -> set[str]:
+    global _DYNAMIC_SOURCES
+    if _DYNAMIC_SOURCES is not None:
+        return _DYNAMIC_SOURCES
+
+    _DYNAMIC_SOURCES = set(torch.compiler.config.dynamic_sources.split(","))
+    return _DYNAMIC_SOURCES
+
+
 # Tracks the sources of all fake tensors we wrap in Dynamo.
 # Used by shape guard computation.
 @dataclasses.dataclass
@@ -2669,6 +2682,7 @@ def _automatic_dynamic(
         unimplemented("torch.compile does not support strided NestedTensor")
 
     name = source.name()
+    dynamic_sources = get_dynamic_sources()
     prior_policy = tx.output.tracing_context.tensor_to_context.get(e, None)
     shape_env_to_source_to_symbol_cache = (
         prior_policy.shape_env_to_source_to_symbol_cache if prior_policy else None
@@ -2807,7 +2821,7 @@ def update_dim2constraint(dim, constraint_range, name):
 
         # Reflect the user directive in the frame_state
         # For dynamic, apply None always
-        if marked_dynamic:
+        if marked_dynamic or source.name() in dynamic_sources:
             # TODO: This can be batched
             # TODO: Doing this here is kind of sus, maybe better to set this
             # up when we initially created the FrameStateSizeEntry to bong
 
@@ -63,5 +63,13 @@
 A common use case for such a tag is to break caches.
 """
 
+dynamic_sources: str = Config(
+    env_name_default="TORCH_COMPILE_DYNAMIC_SOURCES", default=""
+)
+"""
+Comma delimited list of sources that should be marked as dynamic. Primarily useful for large
+models with graph breaks where you need intermediate tensors and ints to be marked dynamic.
+"""
+
 
 install_config_module(sys.modules[__name__])