From 1aa8ccec88a5cac9c59c7be07402b28d55be9b74 Mon Sep 17 00:00:00 2001 From: Albert Cheng Date: Thu, 11 Jun 2026 11:49:31 -0700 Subject: [PATCH] [Autotune] Hardware-aware dispatch cache and accurate L2 flush - Record an environment fingerprint (GPU name, compute capability, CUDA version, target, tilus version) next to the dispatch table and validate it on load. A mismatch -- or a legacy table without a fingerprint -- is ignored and re-tuned, so a table tuned in one environment (e.g. a different GPU sharing a cache directory) is never silently reused. A field can be relaxed to "*" by hand. - Key the fingerprint on the release base version (e.g. 0.2.1) rather than the full SCM version (0.2.1.dev19+g), so development builds off the same release keep sharing the cache instead of invalidating it on every commit. - Benchmarking: flush the L2 using 2x the device's actual L2 cache size (queried from the device) instead of assuming a fixed 128 MiB, with a 128 MiB floor as fallback. Adds CPU unit tests (tests/lang/test_autotune.py) and documents the cache in the autotuning guide. Validated on B300 (sm103a). Tuning budget/seed and heuristic schedule selection are deferred to a follow-up search-control PR (per review feedback) so they can be designed for non-MMA and varied-MMA kernels first. Signed-off-by: Albert Cheng --- docs/source/programming-guides/autotuning.rst | 19 +++ python/tilus/lang/instantiated_script.py | 88 +++++++++++++- python/tilus/utils/bench_utils.py | 18 ++- tests/lang/test_autotune.py | 114 ++++++++++++++++++ 4 files changed, 232 insertions(+), 7 deletions(-) create mode 100644 tests/lang/test_autotune.py diff --git a/docs/source/programming-guides/autotuning.rst b/docs/source/programming-guides/autotuning.rst index d8edc695..ba858dcf 100644 --- a/docs/source/programming-guides/autotuning.rst +++ b/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``), so that development builds off the same release +continue to share a cache instead of invalidating it on every commit. diff --git a/python/tilus/lang/instantiated_script.py b/python/tilus/lang/instantiated_script.py index 0014cffd..b6a2431e 100644 --- a/python/tilus/lang/instantiated_script.py +++ b/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``). 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]) diff --git a/python/tilus/utils/bench_utils.py b/python/tilus/utils/bench_utils.py index b8f1a708..895f1150 100644 --- a/python/tilus/utils/bench_utils.py +++ b/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 diff --git a/tests/lang/test_autotune.py b/tests/lang/test_autotune.py new file mode 100644 index 00000000..624d7a0b --- /dev/null +++ b/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"), 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)