syerva-g · Oct 7, 2026 · Oct 14, 2026
diff --git a/test/higher_order_ops/test_invoke_subgraph.py b/test/higher_order_ops/test_invoke_subgraph.py
@@ -5412,6 +5412,91 @@ def fn(x):
         # second call should reuse the first trace.
         self.assertEqual(count(), 1)
 
+    def test_subgraph_reuse_nested_region_guards(self):
+        """Outer nested_compile_region must inherit guards from inner nested_compile_region."""
+
+        class Layer(torch.nn.Module):
+            def __init__(self, flag: bool):
+                super().__init__()
+                self.flag = flag
+                self.w = torch.nn.Parameter(torch.ones(4, 4))
+
+        @nested_compile_region
+        def inner(layer, x):
+            if layer.flag:
+                return x @ layer.w
+            return x + 1.0
+
+        @nested_compile_region
+        def outer(layer, x):
+            return inner(layer, x)
+
+        class M(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.layers = torch.nn.ModuleList([Layer(True), Layer(False)])
+
+            def forward(self, x):
+                for layer in self.layers:
+                    x = outer(layer, x)
+                return x
+
+        m = M()
+        x = torch.randn(4, 4)
+        ref = m(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(m, backend="aot_eager", fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+        # layer[0] (flag=True) and layer[1] (flag=False) must both trace
+        # both outer and inner (2 * 2 = 4 speculate calls).
+        self.assertEqual(count(), 4)
+
+    def test_subgraph_reuse_nested_region_stamped_out_inner_guards(self):
+        """Outer nested_compile_region must inherit guards even when inner hits reuse cache."""
+
+        class Layer(torch.nn.Module):
+            def __init__(self, flag: bool):
+                super().__init__()
+                self.flag = flag
+                self.w = torch.nn.Parameter(torch.ones(4, 4))
+
+        @nested_compile_region
+        def inner(layer, x):
+            if layer.flag:
+                return x @ layer.w
+            return x + 1.0
+
+        @nested_compile_region
+        def outer(layer, x):
+            return inner(layer, x)
+
+        class M(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.pre_layer = Layer(True)
+                self.layers = torch.nn.ModuleList([Layer(True), Layer(False)])
+
+            def forward(self, x):
+                # Warm up inner's reuse cache before outer is ever traced.
+                x = inner(self.pre_layer, x)
+                for layer in self.layers:
+                    x = outer(layer, x)
+                return x
+
+        m = M()
+        x = torch.randn(4, 4, requires_grad=True)
+        ref = m(x)
+
+        with self._count_speculate_calls() as count:
+            res = torch.compile(m, backend="aot_eager", fullgraph=True)(x)
+
+        self.assertEqual(ref, res)
+        # 1 trace for inner(pre_layer), 1 trace for outer(layers[0]) (which stamps out inner),
+        # and 2 traces for outer(layers[1]) + inner(layers[1]) = 4 total traces.
+        self.assertEqual(count(), 4)
+
 
 @skipIfTorchDynamo("Not a torch._dynamo test")
 @parameterized_class(

@@ -2023,6 +2023,9 @@ def visit(vt: VariableTracker) -> None:
             # - `graph_output_vts`: Only the tensor/symint VTs that are actual
             #   FX graph outputs (basically the vts associated with graph outputs)
             # - `tracing_info`: Properties observed during subgraph tracing
+            if subtracer.parent is not None:
+                subtracer.parent.traced_sources.update(subtracer.traced_sources)
+
             tracing_info = SubgraphTracingInfo(
                 side_effect_stack=subtracer.side_effect_stack,
                 traced_sources=subtracer.traced_sources,

diff --git a/torch/_dynamo/variables/invoke_subgraph.py b/torch/_dynamo/variables/invoke_subgraph.py
@@ -921,6 +921,7 @@ def save_reuse_entry(
     max_reuse_entries: int = 8,
     condition: "InvokeSubgraphReuseCondition | None" = None,
     hash_key: int | None = None,
+    traced_sources: OrderedSet[Source] | None = None,
 ) -> None:
     """Save a traced subgraph into the reuse cache for future cache hits.
 
@@ -977,6 +978,11 @@ def save_reuse_entry(
         if isinstance(t, torch.Tensor)
     ]
 
+    effective_traced_sources = (
+        traced_sources
+        if traced_sources is not None
+        else (condition.traced_sources if condition is not None else OrderedSet())
+    )
     entry = InvokeSubgraphReuseEntry(
         body_name=body_name,
         body_gmod=body_gmod,
@@ -989,6 +995,8 @@ def save_reuse_entry(
         # rewrite captured variable sources for the current invocation.
         arg_sources=fingerprint.arg_sources,
         num_user_outputs=num_user_outputs,
+        condition=condition,
+        traced_sources=effective_traced_sources,
     )
     if condition is not None:
         invoke_subgraph_cache.add_reuse_entry(
@@ -1061,6 +1069,7 @@ def stamp_out_subgraph(
     source replacement before we can look up or create the corresponding
     graph placeholders.
     """
+    from torch._dynamo.guards import install_guard
     from torch._dynamo.variables.builder import VariableBuilder
     from torch._dynamo.variables.higher_order_ops import add_call_function, make_attr
 
@@ -1069,6 +1078,38 @@ def stamp_out_subgraph(
 
     source_replacement = build_source_replacement(cached.arg_sources, new_arg_sources)
 
+    if cached.traced_sources:
+        if source_replacement:
+            remapped_sources = OrderedSet(
+                s.clone(lambda src: source_replacement.get(src, src))
+                for s in cached.traced_sources
+            )
+        else:
+            remapped_sources = cached.traced_sources
+        tx.output.current_tracer.traced_sources.update(remapped_sources)
+
+    if cached.condition is not None and source_replacement:
+        # Re-install the cached reuse condition's guards under the current
+        # call's sources so that an enclosing nested_compile_region picks them
+        # up when it builds its own reuse condition from traced_sources.
+        #
+        # Skip guards on the args themselves and TENSOR_MATCH guards: the
+        # caller already guarded the new args when it wrapped them, and
+        # captured tensor sources are re-wrapped through VariableBuilder below
+        # (which installs their TENSOR_MATCH).  TENSOR_MATCH is also bound to
+        # the original tensor via functools.partial, so cloning it onto a new
+        # source would install a second, stale tensor guard on that source and
+        # fail guard construction with "Duplicate tensors found".
+        old_arg_sources = {s for s in cached.arg_sources if s is not None}
+        for source, _handler, _expected, guard in cached.condition.guards:
+            if source in old_arg_sources or guard.create_fn_name() == "TENSOR_MATCH":
+                continue
+            new_source = source.clone(lambda src: source_replacement.get(src, src))
+            if new_source != source:
+                new_guard = new_source.make_guard(guard.create_fn)
+                new_guard._force_dict_keys_match = guard._force_dict_keys_match
+                install_guard(new_guard)
+
     new_lifted_args = []
     # Shared resolution context so get_value memoizes intermediate results
     # (e.g. L['self'].layers) across all freevars in this stamp-out.
@@ -1470,6 +1511,7 @@ def _call_function(
                     example_value,
                     max_reuse_entries,
                     hash_key=hash_key,  # type: ignore[possibly-undefined]
+                    traced_sources=traced_sources,
                 )
             else:
                 traced_sources = tracing_info.traced_sources
@@ -1494,6 +1536,7 @@ def _call_function(
                             example_value,
                             max_reuse_entries,
                             condition=condition,
+                            traced_sources=traced_sources,
                         )
 
         return _call_function_with_auto_output_flattening(  # type: ignore[return-value]

diff --git a/torch/_guards.py b/torch/_guards.py
@@ -817,6 +817,8 @@ class InvokeSubgraphReuseEntry:
     # The graph may have additional outputs from side-effect intermediates;
     # stamp_out_subgraph uses this to return only the user-visible slice.
     num_user_outputs: int = 0
+    condition: InvokeSubgraphReuseCondition | None = None
+    traced_sources: OrderedSet[Source] = dataclasses.field(default_factory=OrderedSet)
 
 
 @dataclass