Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions helion/autotuner/base_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,8 +157,13 @@ def _get_cache_info_message(self) -> str:
"""Return a message describing where the cache is and how to clear it."""
return ""

def autotune(self) -> Config:
if os.environ.get("HELION_SKIP_CACHE", "") not in {"", "0", "false", "False"}:
def autotune(self, *, skip_cache: bool = False) -> Config:
if skip_cache or os.environ.get("HELION_SKIP_CACHE", "") not in {
"",
"0",
"false",
"False",
}:
return self.autotuner.autotune()

if (config := self.get()) is not None:
Expand Down
4 changes: 2 additions & 2 deletions helion/autotuner/base_search.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ class BaseAutotuner(abc.ABC):
"""

@abc.abstractmethod
def autotune(self) -> Config:
def autotune(self, *, skip_cache: bool = False) -> Config:
raise NotImplementedError


Expand Down Expand Up @@ -420,7 +420,7 @@ def parallel_benchmark(
results.append((config, fn, inf))
return results

def autotune(self) -> Config:
def autotune(self, *, skip_cache: bool = False) -> Config:
"""
Perform autotuning to find the best configuration.

Expand Down
10 changes: 6 additions & 4 deletions helion/runtime/kernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -253,7 +253,7 @@ def autotune(
self,
args: Sequence[object],
*,
force: bool = False,
force: bool = True,
**options: object,
) -> Config:
"""
Expand Down Expand Up @@ -475,7 +475,7 @@ def autotune(
self,
args: Sequence[object],
*,
force: bool = False,
force: bool = True,
**kwargs: object,
) -> Config:
"""
Expand Down Expand Up @@ -508,7 +508,9 @@ def autotune(
config = FiniteSearch(self, args, self.configs).autotune()
else:
self.settings.check_autotuning_disabled()
config = self.settings.autotuner_fn(self, args, **kwargs).autotune()
config = self.settings.autotuner_fn(self, args, **kwargs).autotune(
skip_cache=force
)

self.set_config(config)
return config
Expand Down Expand Up @@ -623,7 +625,7 @@ def __call__(self, *args: object) -> _R:
if (config := self._implicit_config()) is not None:
self.set_config(config)
else:
self.autotune(args)
self.autotune(args, force=False)
assert self._run is not None

assert self._config is not None
Expand Down
Loading