From 9c5c4336ba8b6d79020190d2649bd1afee749a96 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Wed, 29 Jul 2026 21:02:41 -0700 Subject: [PATCH 1/3] Fix CUDA autotuning with released KV storage --- backends/cuda/cuda_backend.py | 4 ++++ backends/cuda/tests/test_cuda_export.py | 19 ++++++++++++++++++- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index bf28077c62c..18a37f124ab 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -122,6 +122,10 @@ 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 _get_const_synthesize_zeros(self, name): diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 2bd3d4db109..71428b9c192 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -8,7 +8,10 @@ from typing import Tuple import torch -from executorch.backends.cuda.cuda_backend import CudaBackend +from executorch.backends.cuda.cuda_backend import ( + CudaBackend, + _compile_time_cpu_clones, +) 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 +20,20 @@ class TestCudaBackendCompileOptions(unittest.TestCase): + def test_low_memory_clone_rehydrates_emptied_tensor_for_autotuning(self): + tensor = torch.empty_strided((2, 3), (3, 1)) + tensor.untyped_storage().resize_(0) + + with _compile_time_cpu_clones(torch.device("cuda")): + from torch._inductor import compile_fx + + clone = compile_fx.clone_preserve_strides(tensor) + + self.assertEqual(clone.shape, tensor.shape) + self.assertEqual(clone.stride(), tensor.stride()) + self.assertEqual(clone.untyped_storage().nbytes(), 6 * tensor.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")] From 6cdd477e3bbed3184a871253b545bfdb915c7e3f Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Wed, 29 Jul 2026 21:10:21 -0700 Subject: [PATCH 2/3] Format CUDA export regression test --- backends/cuda/tests/test_cuda_export.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 71428b9c192..b94cabd2519 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -8,10 +8,7 @@ from typing import Tuple import torch -from executorch.backends.cuda.cuda_backend import ( - CudaBackend, - _compile_time_cpu_clones, -) +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 c86b7c2fdfdd0706f02b3b535c73f126071faecc Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Thu, 30 Jul 2026 09:53:04 -0700 Subject: [PATCH 3/3] Rehydrate all emptied CUDA autotune arguments --- backends/cuda/cuda_backend.py | 17 +++++++++++++++ backends/cuda/tests/test_cuda_export.py | 28 +++++++++++++++++-------- 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 18a37f124ab..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 @@ -128,6 +130,19 @@ def _cpu_clone_preserve_strides(x: torch.Tensor) -> torch.Tensor: 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 @@ -156,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: @@ -166,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 b94cabd2519..0811b926e69 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -5,6 +5,7 @@ # LICENSE file in the root directory of this source tree. import unittest +from types import SimpleNamespace from typing import Tuple import torch @@ -18,18 +19,27 @@ class TestCudaBackendCompileOptions(unittest.TestCase): def test_low_memory_clone_rehydrates_emptied_tensor_for_autotuning(self): - tensor = torch.empty_strided((2, 3), (3, 1)) - tensor.untyped_storage().resize_(0) + 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) - with _compile_time_cpu_clones(torch.device("cuda")): - from torch._inductor import compile_fx + 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"] - clone = compile_fx.clone_preserve_strides(tensor) + with _compile_time_cpu_clones(torch.device("cuda")): + cloned_args, _ = autotuner.maybe_clone_args(set(), read_only, mutated) - self.assertEqual(clone.shape, tensor.shape) - self.assertEqual(clone.stride(), tensor.stride()) - self.assertEqual(clone.untyped_storage().nbytes(), 6 * tensor.element_size()) - self.assertEqual(torch.count_nonzero(clone), 0) + 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(