diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index bf28077c62c..9b1bec49576 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -89,11 +89,13 @@ def _compile_time_cpu_clones(target_device: torch.device): from torch._inductor import compile_fx as _cfx, graph as _graph from torch._inductor.codegen.cpp_wrapper_cpu import CppWrapperCpu as _Cpp from torch._inductor.graph import GraphLowering as _GL + from torch._inductor.runtime.triton_heuristics import CachingAutotuner as _Autotuner orig_clone = _cfx.clone_preserve_strides orig_codegen_device = _Cpp.codegen_device orig_get_const = _GL.get_original_value_of_constant orig_is_same = _graph.is_same_tensor + orig_maybe_clone_args = _Autotuner.maybe_clone_args def _is_same_skip_emptied(data, value): # KV buffers freed via resize_(0) all have data_ptr 0, so the stock @@ -122,8 +124,25 @@ def _cpu_clone_preserve_strides(x: torch.Tensor) -> torch.Tensor: if _is_emptied(x): return _full_zeros_preserving_strides(x, "cpu") return orig_clone(x).cpu() + if _is_emptied(x): + # Autotuning needs device-backed inputs. Rehydrate the zero-filled + # KV buffer temporarily instead of cloning its freed storage. + return _full_zeros_preserving_strides(x, x.device) return orig_clone(x) + def _maybe_clone_args_with_rehydrated_kv(self, exclude, *args, **kwargs): + # Inductor only clones mutated arguments before benchmarking. A fused + # kernel can also receive read-only views of an emptied KV buffer, so + # rehydrate every emptied tensor argument before the clone step. + def rehydrate(arg): + if _is_emptied(arg): + return _full_zeros_preserving_strides(arg, arg.device) + return arg + + args = tuple(rehydrate(arg) for arg in args) + kwargs = {name: rehydrate(arg) for name, arg in kwargs.items()} + return orig_maybe_clone_args(self, exclude, *args, **kwargs) + 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 @@ -152,6 +171,7 @@ def _codegen_device_target_aware(self, device): _Cpp.codegen_device = _codegen_device_target_aware _GL.get_original_value_of_constant = _get_const_synthesize_zeros _graph.is_same_tensor = _is_same_skip_emptied + _Autotuner.maybe_clone_args = _maybe_clone_args_with_rehydrated_kv prev_active = getattr(_CPU_CLONE_GUARD, "active", False) _CPU_CLONE_GUARD.active = True try: @@ -162,6 +182,7 @@ def _codegen_device_target_aware(self, device): _Cpp.codegen_device = orig_codegen_device _GL.get_original_value_of_constant = orig_get_const _graph.is_same_tensor = orig_is_same + _Autotuner.maybe_clone_args = orig_maybe_clone_args 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 2bd3d4db109..0811b926e69 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -5,10 +5,11 @@ # LICENSE file in the root directory of this source tree. import unittest +from types import SimpleNamespace from typing import Tuple 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 @@ -17,6 +18,29 @@ class TestCudaBackendCompileOptions(unittest.TestCase): + def test_low_memory_clone_rehydrates_emptied_tensor_for_autotuning(self): + read_only = torch.empty_strided((2, 3), (3, 1)) + mutated = torch.empty_strided((2, 3), (3, 1)) + read_only.untyped_storage().resize_(0) + mutated.untyped_storage().resize_(0) + + from torch._inductor.runtime.triton_heuristics import CachingAutotuner + + autotuner = object.__new__(CachingAutotuner) + autotuner.fn = SimpleNamespace(arg_names=["read_only", "mutated"]) + autotuner.mutated_arg_names = ["mutated"] + + with _compile_time_cpu_clones(torch.device("cuda")): + cloned_args, _ = autotuner.maybe_clone_args(set(), read_only, mutated) + + for clone in cloned_args: + self.assertEqual(clone.shape, read_only.shape) + self.assertEqual(clone.stride(), read_only.stride()) + self.assertEqual( + clone.untyped_storage().nbytes(), 6 * read_only.element_size() + ) + self.assertEqual(torch.count_nonzero(clone), 0) + def test_emulate_precision_casts_compile_spec(self): options = CudaBackend.get_aoti_compile_options( [CompileSpec(key="emulate_precision_casts", value=b"OFF")]