[Autotune] Hardware-aware dispatch cache and accurate L2 flush

docs/source/programming-guides/autotuning.rst

@@ -68,3 +68,22 @@ When we launch the kernel, tilus will automatically compile the kernel with all
 The kernels will be compiled in parallel when we first call the kernel with a specific input size triggering the JIT
 compilation (:doc:`tilus-script`). We can use :py:func:`tilus.option.parallel_workers` to control the number of
 parallel workers to compile the kernels.
+
+Hardware-aware tuning cache
+---------------------------
+
+The best schedule for a given input is hardware- and toolchain-specific, so the tuning result (the
+*dispatch table* that maps an input bucket to the winning schedule) must not be reused across
+different environments. The compiled kernels are already keyed by target architecture, and the
+dispatch table additionally records an environment fingerprint -- the GPU name, compute capability,
+CUDA version, target, and tilus version -- next to it.
+
+When a dispatch table is loaded, this fingerprint is compared against the current environment. If any
+field differs (for example, a table tuned on a B200 being picked up on a B300 through a shared cache
+directory), the table is ignored and the kernel is re-tuned for the current environment instead of
+silently using a schedule that was optimal elsewhere. Advanced users can relax an individual check by
+editing the ``_metadata`` block in ``dispatch_table.json`` and setting a field to ``"*"``.
+
+The tilus version recorded in the fingerprint is the release base version (e.g. ``0.2.1``) rather than
+the full development version (``0.2.1.dev19+g<hash>``), so that development builds off the same release
+continue to share a cache instead of invalidating it on every commit.

python/tilus/lang/instantiated_script.py

@@ -38,7 +38,7 @@
 from tilus.ir.prog import Program
 from tilus.lang.script import Script
 from tilus.runtime import CompiledProgram, load_compiled_program
-from tilus.target import lazy_init
+from tilus.target import get_current_target, lazy_init
 from tilus.utils import benchmark_func, relative_to_with_walk_up, to_snake_case
 from tilus.utils.multiprocess import parallel_imap
 
@@ -159,6 +159,74 @@ def generate_schedules(
     return schedules
 
 
+def _safe_str(fn: Any) -> str:
+    """Call ``fn`` and stringify its result, returning ``"unknown"`` on any failure."""
+    try:
+        return str(fn())
+    except Exception:
+        return "unknown"
+
+
+def _tilus_version() -> str:
+    """Release-level tilus version used in the cache fingerprint.
+
+    We deliberately key on the *base* release (e.g. ``0.2.1``) rather than the full SCM version
+    (``0.2.1.dev19+g<hash>``). The SCM version changes on every commit during development, which would
+    needlessly invalidate the dispatch cache each time; keying on the base release lets all dev builds
+    off the same release share the cache while distinct releases stay separated.
+    """
+    try:
+        import importlib.metadata as importlib_metadata
+
+        raw = importlib_metadata.version("tilus")
+    except Exception:
+        raw = str(getattr(tilus, "__version__", "unknown"))
+    try:
+        from packaging.version import Version
+
+        return Version(raw).base_version
+    except Exception:
+        # strip a PEP 440 dev / local suffix by hand if packaging is unavailable
+        return raw.split(".dev")[0].split("+")[0]
+
+
+def collect_tuning_metadata() -> dict[str, str]:
+    """Fingerprint of the environment that produced a tuning dispatch table.
+
+    The fastest schedule for a tuning key depends on the GPU and the toolchain, so a dispatch table
+    tuned in one environment must not be silently reused in another (for example, a table tuned on a
+    B200 picked up on a B300 through a shared cache directory). These fields are written next to the
+    dispatch table and compared on load; a mismatch causes the on-disk table to be ignored and
+    re-tuned. This mirrors the metadata-validation approach used by FlashInfer's autotuner cache.
+    """
+    return {
+        "tilus_version": _tilus_version(),
+        "target": _safe_str(get_current_target),
+        "gpu": _safe_str(lambda: torch.cuda.get_device_name(torch.cuda.current_device())),
+        "compute_capability": _safe_str(lambda: "{}.{}".format(*torch.cuda.get_device_capability())),
+        "cuda_version": _safe_str(lambda: torch.version.cuda),
+    }
+
+
+def tuning_metadata_matches(saved: Any, current: Mapping[str, str]) -> bool:
+    """Return whether a saved metadata block is compatible with the current environment.
+
+    Every field present in ``current`` must match the corresponding field in ``saved``. A saved field
+    set to the wildcard ``"*"`` matches any value, allowing advanced users to relax individual checks
+    by editing the cache file by hand (as in FlashInfer). A missing or non-mapping ``saved`` (e.g. a
+    legacy cache file without metadata) never matches.
+    """
+    if not isinstance(saved, dict):
+        return False
+    for key, current_value in current.items():
+        saved_value = saved.get(key)
+        if saved_value == "*":
+            continue
+        if saved_value != current_value:
+            return False
+    return True
+
+
 class CallParameters:
     """
     Analyze the parameters in the Script '__call__' function.
@@ -704,17 +772,25 @@ def _pick_best_program(self, args: Sequence[Any]) -> CompiledProgram:
 
     def load_dispatch_table(self):
         table_path = self.cache_dir / "dispatch_table.json"
-        if table_path.exists():
-            with open(table_path, "r") as f:
-                entries = json.load(f)
-            self.dispatch_table = {tuple(key): value for key, value in entries}
+        if not table_path.exists():
+            return
+        with open(table_path, "r") as f:
+            data = json.load(f)
+        # The fastest schedule for a tuning key is environment-specific, so a dispatch table is only
+        # reused when its saved environment fingerprint matches the current one. Otherwise (including a
+        # legacy table that predates the fingerprint) it is ignored and the kernel is re-tuned.
+        saved_meta = data.get("_metadata") if isinstance(data, dict) else None
+        if not tuning_metadata_matches(saved_meta, collect_tuning_metadata()):
+            return
+        self.dispatch_table = {tuple(key): value for key, value in data["entries"]}
 
     def dump_dispatch_table(self):
         table_path = self.cache_dir / "dispatch_table.json"
         table_txt_path = self.cache_dir / "dispatch_table.txt"
         entries = [[list(key), value] for key, value in self.dispatch_table.items()]
+        data = {"_metadata": collect_tuning_metadata(), "entries": entries}
         with open(table_path, "w") as f:
-            json.dump(entries, f)
+            json.dump(data, f)
         headers = []
         for idx in self.call_params.tuning_params:
             headers.append(self.call_params.param_names[idx])

python/tilus/utils/bench_utils.py

@@ -51,13 +51,29 @@ def cuda_sleep(nanoseconds: int) -> None:
     torch.cuda._sleep(nanoseconds)
 
 
+@functools.lru_cache(maxsize=None)
+def _l2_clear_nbytes(device_index: int) -> int:
+    """Return the number of bytes to write in order to evict the L2 cache on the given device.
+
+    Writing a slab twice the device's L2 cache size reliably flushes it, so that each benchmarked
+    iteration starts from a cold L2. The actual L2 size is queried from the device (it varies across
+    architectures); a 128 MiB floor preserves the previous behavior when the size cannot be determined.
+    """
+    floor = 128 * 1024 * 1024
+    try:
+        l2_size = torch.cuda.get_device_properties(device_index).L2_cache_size
+    except Exception:
+        l2_size = 0
+    return max(2 * int(l2_size), floor)
+
+
 def benchmark_func(
     run_func: Callable[[], Any],
     warmup: int = 1,
     repeat: int = 5,
     clear_l2_cache: bool = True,
 ) -> float:
-    num_bytes = 128 * 1024 * 1024
+    num_bytes = _l2_clear_nbytes(torch.cuda.current_device())
     memory_slab = torch.empty(num_bytes, dtype=torch.int8, device="cuda")
 
     assert repeat >= 1

tests/lang/test_autotune.py

@@ -0,0 +1,114 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for the autotune schedule-space generation and the hardware-aware dispatch cache.
+
+These tests exercise pure-Python autotuner logic and do not require a GPU.
+"""
+
+from typing import Any, Sequence
+
+import tilus
+from tilus import float16, int32
+from tilus.lang.instantiated_script import (
+    _tilus_version,
+    collect_tuning_metadata,
+    generate_schedules,
+    span_space,
+    tuning_metadata_matches,
+)
+
+
+def test_span_space_single_and_grouped_keys():
+    space: dict[str, Sequence[Any]] = {"m": [1, 2, 3], "n, k": [[1, 2], [3, 4]]}
+    spanned = span_space(space)
+    # 3 choices for m, 2 choices for (n, k) -> 6 combinations
+    assert len(spanned) == 6
+    assert {"m": 1, "n": 1, "k": 2} in spanned
+    assert {"m": 3, "n": 3, "k": 4} in spanned
+    # every spanned schedule has the flattened keys
+    for entry in spanned:
+        assert set(entry) == {"m", "n", "k"}
+
+
+@tilus.autotune("block_m, block_n", [[64, 64], [128, 64], [128, 128]])
+@tilus.autotune("block_k", [16, 32, 64])
+class _DummyMatmul(tilus.Script):
+    def __init__(self, block_m: int, block_n: int, block_k: int):
+        super().__init__()
+        self.block_m = block_m
+        self.block_n = block_n
+        self.block_k = block_k
+
+    def __call__(self, m: int32, a_ptr: ~float16):  # pragma: no cover - never executed
+        pass
+
+
+def _dummy_schedules():
+    space = getattr(_DummyMatmul, "_autotune_space")
+    return generate_schedules(space, _DummyMatmul, script_args=(), script_kwargs={})
+
+
+def test_generate_schedules_cartesian_product():
+    schedules = _dummy_schedules()
+    # 3 (block_m, block_n) x 3 (block_k) = 9 schedules
+    assert len(schedules) == 9
+    assert {"block_m": 64, "block_n": 64, "block_k": 16} in schedules
+    assert {"block_m": 128, "block_n": 128, "block_k": 64} in schedules
+
+
+# ---------------------------------------------------------------------------
+# Hardware-aware dispatch-cache fingerprint
+# ---------------------------------------------------------------------------
+
+
+def test_collect_tuning_metadata_has_expected_keys():
+    meta = collect_tuning_metadata()
+    assert set(meta) == {"tilus_version", "target", "gpu", "compute_capability", "cuda_version"}
+    # all values must be strings (never None) so they serialize and compare cleanly
+    assert all(isinstance(v, str) for v in meta.values())
+
+
+def test_tuning_version_is_release_base():
+    # The fingerprint must key on the release base version (e.g. "0.2.1"), not the full SCM/dev
+    # version ("0.2.1.dev19+g<hash>"), so dev builds off the same release keep sharing the cache.
+    version = _tilus_version()
+    assert ".dev" not in version
+    assert "+" not in version
+
+
+def test_tuning_metadata_matches_identical():
+    meta = {"gpu": "NVIDIA B300", "compute_capability": "10.3", "cuda_version": "13.0"}
+    assert tuning_metadata_matches(meta, meta)
+
+
+def test_tuning_metadata_matches_detects_gpu_mismatch():
+    saved = {"gpu": "NVIDIA B200", "compute_capability": "10.0"}
+    current = {"gpu": "NVIDIA B300", "compute_capability": "10.3"}
+    assert not tuning_metadata_matches(saved, current)
+
+
+def test_tuning_metadata_matches_wildcard():
+    saved = {"gpu": "*", "compute_capability": "10.3"}
+    current = {"gpu": "NVIDIA B300", "compute_capability": "10.3"}
+    assert tuning_metadata_matches(saved, current)
+
+
+def test_tuning_metadata_matches_rejects_legacy_or_missing():
+    current = {"gpu": "NVIDIA B300"}
+    # legacy cache files without a metadata mapping must never match
+    assert not tuning_metadata_matches(None, current)
+    assert not tuning_metadata_matches([], current)
+    # a metadata block missing a required field does not match
+    assert not tuning_metadata_matches({}, current)