From 70e05437a1a6da5cd8aeb5950b6667077d374ffa Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Mon, 12 Oct 2026 22:21:21 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/cuda/cuda_backend.py | 18 ++++ backends/cuda/tests/test_cuda_export.py | 122 +++++++++++++++++++++++- 2 files changed, 138 insertions(+), 2 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index f28faabb624..2f2c66b2c08 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -260,6 +260,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 @@ -304,6 +305,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 @@ -337,6 +353,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: @@ -350,6 +367,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: diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 8329f42e522..6cd4d909675 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -13,11 +13,15 @@ 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 +39,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 +69,97 @@ 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 +340,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