Fix CUDA autotuning with released KV storage

backends/cuda/cuda_backend.py

@@ -83,17 +83,19 @@
 
 
 @contextlib.contextmanager
 def _compile_time_cpu_clones(target_device: torch.device):
     """Force AOTI's mutated-buffer clones onto CPU while preserving the
     serialized constants' target 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 @@
             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 @@
     _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 @@
         _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:

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")]