[executorch][cuda] Retain emptied KV backing across JIT autotuning

backends/cuda/cuda_backend.py

@@ -261,6 +261,7 @@ def _force_external_weights_for_fqn_binding(consts_size):
         return orig_determine_aoti_mmap_flags(consts_size)
 
     orig_autotuner_run = _Autotuner.run
+    orig_jit_autotune = _GL._run_jit_variant_for_autotune
 
     def _is_same_skip_emptied(data, value):
         # KV buffers freed via resize_(0) all have data_ptr 0, so the stock
@@ -305,6 +306,21 @@ def _autotuner_run_with_rehydrated_emptied_args(self, *args, **kwargs):
         with _rehydrate_emptied_tensors(tensors):
             return orig_autotuner_run(self, *args, **kwargs)
 
+    def _jit_autotune_with_rehydrated_constants(
+        self, wrapper_code, kernel_code, extract_real_inputs, kernel_names
+    ):
+        # The OFF path launches C++ kernels directly, bypassing Autotuner.run.
+        # Restore backing before pointer extraction and keep it through completion.
+        backing = (
+            _rehydrate_emptied_tensors(self.constants.values())
+            if _is_cpu_clone_active()
+            else contextlib.nullcontext()
+        )
+        with backing:
+            return orig_jit_autotune(
+                self, wrapper_code, kernel_code, extract_real_inputs, kernel_names
+            )
+
     def _get_const_synthesize_zeros(self, name):
         # AOTI serializes each constant via get_original_value_of_constant ->
         # _to_bytes. For KV buffers we freed with resize_(0) this would otherwise
@@ -338,6 +354,7 @@ def _codegen_device_target_aware(self, device):
     )
     _codecache.determine_aoti_mmap_flags = _force_external_weights_for_fqn_binding
     _Autotuner.run = _autotuner_run_with_rehydrated_emptied_args
+    _GL._run_jit_variant_for_autotune = _jit_autotune_with_rehydrated_constants
     prev_active = getattr(_CPU_CLONE_GUARD, "active", False)
     _CPU_CLONE_GUARD.active = True
     try:
@@ -351,6 +368,7 @@ def _codegen_device_target_aware(self, device):
         _codecache.TensorProperties = orig_tensor_properties
         _codecache.determine_aoti_mmap_flags = orig_determine_aoti_mmap_flags
         _Autotuner.run = orig_autotuner_run
+        _GL._run_jit_variant_for_autotune = orig_jit_autotune
 
 
 def _is_kv_buffer(name, v) -> bool:

backends/cuda/tests/test_cuda_export.py

@@ -13,11 +13,12 @@
 from unittest.mock import patch
 
 import torch
-from executorch.backends.cuda.cuda_backend import CudaBackend
+from executorch.backends.cuda.cuda_backend import _compile_time_cpu_clones, CudaBackend
 from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
 from executorch.examples.models.toy_model import SdpaModule
 from executorch.exir import EdgeCompileConfig, schema, to_edge_transform_and_lower
 from executorch.exir.backend.compile_spec_schema import CompileSpec
+from torch._inductor.graph import GraphLowering
 from torch.export import export
 
 
@@ -35,7 +36,6 @@ def test_low_memory_triton_reduction_loads_stay_loop_scoped(self):
         self.assertIs(IndexingOptions.has_rmask, original_has_rmask)
 
     def test_low_memory_autotune_rehydrates_aliased_storage_for_full_run(self):
-        from executorch.backends.cuda.cuda_backend import _compile_time_cpu_clones
         from torch._inductor.runtime.triton_heuristics import CachingAutotuner
 
         base = torch.empty_strided((4, 4), (4, 1))
@@ -66,6 +66,99 @@ def fake_run(_autotuner, base_arg, view_arg, *, stream):
         self.assertEqual(result, "run-result")
         self.assertEqual(base.untyped_storage().nbytes(), 0)
 
+    def test_low_memory_jit_rehydrates_aliases_before_input_extraction(self):
+        base = torch.empty_strided((4, 4), (4, 1))
+        view = base[:, 1:]
+        storage = base.untyped_storage()
+        storage.resize_(0)
+        graph = object.__new__(GraphLowering)
+        weight = torch.ones(3)
+        weight_ptr = weight.data_ptr()
+        graph.constants = {"base": base, "view": view, "weight": weight}
+        original_run = GraphLowering._run_jit_variant_for_autotune
+
+        def extract_real_inputs():
+            self.assertEqual(storage.nbytes(), 64)
+            self.assertEqual(torch.count_nonzero(base), 0)
+            return base, view
+
+        def jit_run(_graph, wrapper, kernel, extract, names):
+            self.assertIs(_graph, graph)
+            self.assertEqual((wrapper, kernel, names), ("wrapper", "kernel", ["k"]))
+            base_arg, view_arg = extract()
+            self.assertIs(base_arg, base)
+            self.assertIs(view_arg, view)
+            self.assertEqual(view_arg.untyped_storage()._cdata, storage._cdata)
+            self.assertEqual(view_arg.storage_offset(), 1)
+            base_arg.fill_(2)
+            torch.testing.assert_close(view_arg, torch.full((4, 3), 2.0))
+            self.assertEqual(weight.data_ptr(), weight_ptr)
+            torch.testing.assert_close(weight, torch.ones(3))
+            return "jit-result"
+
+        with mock.patch.object(GraphLowering, "_run_jit_variant_for_autotune", jit_run):
+            with _compile_time_cpu_clones(torch.device("cuda")):
+                result = graph._run_jit_variant_for_autotune(
+                    "wrapper", "kernel", extract_real_inputs, ["k"]
+                )
+            self.assertIs(GraphLowering._run_jit_variant_for_autotune, jit_run)
+        self.assertIs(GraphLowering._run_jit_variant_for_autotune, original_run)
+        self.assertEqual(result, "jit-result")
+        self.assertEqual(storage.nbytes(), 0)
+        self.assertEqual(
+            (base.shape, base.stride(), view.storage_offset()), ((4, 4), (4, 1), 1)
+        )
+
+    def test_low_memory_jit_failure_releases_backing_and_restores_hook(self):
+        tensor = torch.empty(4)
+        storage = tensor.untyped_storage()
+        storage.resize_(0)
+        graph = object.__new__(GraphLowering)
+        graph.constants = {"kv": tensor}
+        original_run = GraphLowering._run_jit_variant_for_autotune
+
+        def jit_run(_graph, _wrapper, _kernel, extract, _names):
+            self.assertGreater(storage.nbytes(), 0)
+            self.assertNotEqual(extract(), 0)
+            raise RuntimeError("JIT execution failed")
+
+        with mock.patch.object(GraphLowering, "_run_jit_variant_for_autotune", jit_run):
+            with self.assertRaisesRegex(RuntimeError, "JIT execution failed"):
+                with _compile_time_cpu_clones(torch.device("cuda")):
+                    graph._run_jit_variant_for_autotune(
+                        "wrapper", "kernel", tensor.data_ptr, ["k"]
+                    )
+            self.assertIs(GraphLowering._run_jit_variant_for_autotune, jit_run)
+        self.assertIs(GraphLowering._run_jit_variant_for_autotune, original_run)
+        self.assertEqual(storage.nbytes(), 0)
+
+    def test_low_memory_jit_does_not_rehydrate_on_another_thread(self):
+        tensor = torch.empty(4)
+        storage = tensor.untyped_storage()
+        storage.resize_(0)
+        graph = object.__new__(GraphLowering)
+        graph.constants = {"kv": tensor}
+        observed = []
+
+        def jit_run(_graph, _wrapper, _kernel, extract, _names):
+            return storage.nbytes(), extract()
+
+        def run():
+            observed.append(
+                graph._run_jit_variant_for_autotune(
+                    "wrapper", "kernel", lambda: "unchanged", ["k"]
+                )
+            )
+
+        with mock.patch.object(GraphLowering, "_run_jit_variant_for_autotune", jit_run):
+            with _compile_time_cpu_clones(torch.device("cuda")):
+                worker = threading.Thread(target=run)
+                worker.start()
+                worker.join(timeout=5)
+                self.assertFalse(worker.is_alive())
+        self.assertEqual(observed, [(0, "unchanged")])
+        self.assertEqual(storage.nbytes(), 0)
+
     def test_emulate_precision_casts_compile_spec(self):
         options = CudaBackend.get_aoti_compile_options(
             [CompileSpec(key="emulate_precision_casts", value=b"OFF")]
@@ -246,6 +339,30 @@ def setUp(self):
         if not torch.cuda.is_available():
             self.skipTest("CUDA is not available")
 
+    def test_low_memory_jit_synchronizes_cuda_backing_before_release(self):
+        tensor = torch.empty(16, device="cuda")
+        storage = tensor.untyped_storage()
+        storage.resize_(0)
+        graph = object.__new__(GraphLowering)
+        graph.constants = {"kv": tensor}
+
+        def jit_run(_graph, _wrapper, _kernel, extract, _names):
+            self.assertNotEqual(extract(), 0)
+            tensor.fill_(3)
+            return "cuda-jit-result"
+
+        with mock.patch.object(GraphLowering, "_run_jit_variant_for_autotune", jit_run):
+            with mock.patch.object(
+                torch.cuda, "synchronize", wraps=torch.cuda.synchronize
+            ) as synchronize:
+                with _compile_time_cpu_clones(torch.device("cuda")):
+                    result = graph._run_jit_variant_for_autotune(
+                        "wrapper", "kernel", tensor.data_ptr, ["k"]
+                    )
+                synchronize.assert_called_once_with(tensor.device)
+        self.assertEqual(result, "cuda-jit-result")
+        self.assertEqual(storage.nbytes(), 0)
+
     def test_rehydrated_storage_is_synchronized_before_release(self):
         from executorch.backends.cuda.cuda_backend import _rehydrate_emptied_tensors