From bc9515febdeeac6fa7b339d322adebd3add7e07c Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Wed, 29 Jul 2026 14:08:14 +0800 Subject: [PATCH] [BugFix] Prevent autotuner cache reuse across different outputs and validation settings (#2793) * [BugFix] Include behavioral args in autotune cache keys * [BugFix] Bypass autotune cache for profile callbacks --- .../autotune/test_autotune_cache_key.py | 48 +++++++++++++++++++ tilelang/autotuner/param.py | 3 ++ tilelang/autotuner/tuner.py | 26 ++++++++-- 3 files changed, 72 insertions(+), 5 deletions(-) create mode 100644 testing/python/autotune/test_autotune_cache_key.py diff --git a/testing/python/autotune/test_autotune_cache_key.py b/testing/python/autotune/test_autotune_cache_key.py new file mode 100644 index 000000000..17b8000bb --- /dev/null +++ b/testing/python/autotune/test_autotune_cache_key.py @@ -0,0 +1,48 @@ +"""Tests for auto-tuner cache-key identity.""" + +import inspect +import threading +from dataclasses import replace + +from tilelang.autotuner import AutoTuner +from tilelang.autotuner.param import CompileArgs, ProfileArgs + + +def _kernel(block_size=128): + return block_size + + +def _cache_key(compile_args=None, profile_args=None): + tuner = AutoTuner(_kernel, configs=[{"block_size": 128}]) + tuner.compile_args = compile_args or CompileArgs() + tuner.profile_args = profile_args or ProfileArgs() + return tuner.generate_cache_key(inspect.signature(_kernel).parameters, {}) + + +def test_cache_key_includes_output_indices(): + base_args = CompileArgs(out_idx=[0]) + + assert _cache_key(compile_args=base_args) != _cache_key(compile_args=replace(base_args, out_idx=[1])) + + +def test_cache_key_includes_profile_validation_and_input_behavior(): + base_args = ProfileArgs(skip_check=False, cache_input_tensors=False) + + variants = ( + replace(base_args, skip_check=True), + replace(base_args, cache_input_tensors=True), + ) + + base_key = _cache_key(profile_args=base_args) + assert all(base_key != _cache_key(profile_args=variant) for variant in variants) + + +def test_cache_key_is_disabled_for_profile_callbacks(): + lock = threading.Lock() + + def callback(value): + with lock: + return value + + for callback_field in ("ref_prog", "supply_prog", "manual_check_prog"): + assert _cache_key(profile_args=ProfileArgs(**{callback_field: callback})) is None diff --git a/tilelang/autotuner/param.py b/tilelang/autotuner/param.py index 73afb26b6..f20179313 100644 --- a/tilelang/autotuner/param.py +++ b/tilelang/autotuner/param.py @@ -76,6 +76,7 @@ class CompileArgs: def __hash__(self): """Return a stable hash for cache key construction.""" data = { + "out_idx": self.out_idx, "execution_backend": self.execution_backend, "target": str(self.target), "target_host": str(self.target_host) if self.target_host else None, @@ -136,6 +137,8 @@ class ProfileArgs: "rtol": self.rtol, "atol": self.atol, "max_mismatched_ratio": self.max_mismatched_ratio, + "skip_check": self.skip_check, + "cache_input_tensors": self.cache_input_tensors, } hash_obj = hashlib.sha256(json.dumps(data, sort_keys=True).encode("utf-8")) return int.from_bytes(hash_obj.digest(), byteorder="big") diff --git a/tilelang/autotuner/tuner.py b/tilelang/autotuner/tuner.py index 0b7a3c19d..74fbd8590 100644 --- a/tilelang/autotuner/tuner.py +++ b/tilelang/autotuner/tuner.py @@ -454,9 +454,23 @@ class AutoTuner: "Please provide concrete inputs with `with set_autotune_inputs(...)`." ) - def generate_cache_key(self, parameters: dict[str, Any], extra_parameters: dict[str, Any]) -> AutotuneResult | None: + def generate_cache_key(self, parameters: dict[str, Any], extra_parameters: dict[str, Any]) -> str | None: """Generate a cache key for the auto-tuning process.""" + # Arbitrary callbacks do not have a reliable persistent identity: + # importable functions may serialize by name, while closures may hold + # unpicklable runtime state. Avoid cache reuse rather than risk serving + # results produced with different input or validation behavior. + if any( + callback is not None + for callback in ( + self.profile_args.ref_prog, + self.profile_args.supply_prog, + self.profile_args.manual_check_prog, + ) + ): + return None + # extract parameters from the function signature op_parameters = [] for _, default_value in parameters.items(): @@ -997,7 +1011,7 @@ class AutoTuner: key = self.generate_cache_key(parameters, extra_parameters) with self._lock: - if env.is_cache_enabled() and not env.is_autotune_cache_disabled(): + if key is not None and env.is_cache_enabled() and not env.is_autotune_cache_disabled(): # First check in-memory cache if key in self._memory_cache: # Include PrimFunc name when hitting autotuner memory cache @@ -1077,7 +1091,8 @@ class AutoTuner: # compile the kernel with the provided parameters jit_kernel = self.jit_compile() autotuner_result = AutotuneResult(libcode=jit_kernel.get_kernel_source(), func=jit_kernel.prim_func, kernel=jit_kernel) - self._memory_cache[key] = autotuner_result + if key is not None: + self._memory_cache[key] = autotuner_result return autotuner_result # After confirming tuning will actually run, validate that scalar @@ -1283,10 +1298,11 @@ class AutoTuner: logger.warning("DLPack backend does not support cache saving to disk.") else: with self._lock: - if env.is_cache_enabled() and not env.is_autotune_cache_disabled(): + if key is not None and env.is_cache_enabled() and not env.is_autotune_cache_disabled(): self._save_result_to_disk(key, autotuner_result) - self._memory_cache[key] = autotuner_result + if key is not None: + self._memory_cache[key] = autotuner_result return autotuner_result