[RFC] Implement basic on disk caching #336
Conversation
stack-info: PR: #336, branch: oulgen/stack/26
meta-cla
bot
added
the
CLA Signed
stack-info: PR: #336, branch: oulgen/stack/26
stack-info: PR: #336, branch: oulgen/stack/26
helion/autotuner/cache.py
Outdated
| helion_key: Hash of source code of Helion | ||
| torch_key: Hash of source code of PyTorch | ||
| system_hash: Hash of system information, | ||
| including Triton, current device, cuda/rocm arch version |
There was a problem hiding this comment.
Do we want to include these by default?
I was thinking for this use case, where autotuning is very expensive and caching is off by default, it may make sense to default to a more minimal key.
There was a problem hiding this comment.
My take is correctness is utmost important, and by default we should have as strict key as possible because if we are writing to the file system, people update their helion/triton/torch version, we shouldnt give them bad autotuning results.
Having said this, on L26, one of the TODOs is giving the users customization over the cache key. With that customization, we can allow users to choose to omit one or more of these by-default-included-strict requirements.
What do you think?
There was a problem hiding this comment.
I think baseline here is users copy and pasting the config into their code. This makes the config a BC surface where we need to make sure old configs work with new versions of Helion/PyTorch.
I'm a lot more worried about surprise re-autotuning.
helion/autotuner/cache.py
Outdated
| return hashlib.sha256(repr(self).encode("utf-8")).hexdigest() | ||
|
|
||
|
|
||
| class AutotuneCache: |
helion/autotuner/cache.py
Outdated
| for idx, a in enumerate(args): | ||
| if isinstance(a, torch.Tensor): | ||
| input_dtypes.append((idx, a.dtype)) | ||
| input_shapes.append((idx, a.shape)) |
There was a problem hiding this comment.
Maybe we should mimic/reuse some of the key generation infra in kernel.py (for the in-memory cache). IMO the memory key should be a subset of the disk key. I think right now the memory key includes some extra stuff.
helion/autotuner/cache.py
Outdated
| cache_dir, # pyright: ignore[reportPrivateImportUsage] | ||
| ) | ||
|
|
||
| return Path(cache_dir()) / "helion" / f"{self._get_cache_key()}.best_config" |
There was a problem hiding this comment.
It might be useful to let users specify a filename for the cache so they can more easily distribute it to production environments.
There was a problem hiding this comment.
TODO on L27 is about importing and exporting a file with all the autotune settings. Do you think that would be good enough for should I add an env option to control cache directory?
FWIW we can just do both
helion/autotuner/cache.py
Outdated
|
|
||
| def put(self, config: Config) -> None: | ||
| path = self._get_local_cache_path() | ||
| config.save(path) |
There was a problem hiding this comment.
Should we do the save to a temp file and rename trick here so the change is atomic. I thinking about multi-process races.
test/test_cache.py
Outdated
|
|
||
|
|
||
| class TestCache(TestCase): | ||
| @mock.patch("helion.autotuner.DifferentialEvolutionSearch", new=BasicSearch) |
There was a problem hiding this comment.
Should we add a cleaner way for people to set the autotuner to use?
There was a problem hiding this comment.
Yea, lets define an API that people can set. Out of the scope of this PR but I can look into it. Filed #337
stack-info: PR: #336, branch: oulgen/stack/26
stack-info: PR: #336, branch: oulgen/stack/26
stack-info: PR: #336, branch: oulgen/stack/26
helion/autotuner/base_cache.py
Outdated
|
|
||
| @functools.cache | ||
| def helion_key() -> str: | ||
| from torch._inductor.codecache import build_code_hash |
There was a problem hiding this comment.
inductor imports at top of file throughout PR
| def _get_bound_kernel_cache_key( | ||
| self, args: tuple[object, ...], signature: tuple[Hashable, ...] | ||
| ) -> BoundKernelInMemoryCacheKey | None: | ||
| from ..autotuner.base_cache import BoundKernelInMemoryCacheKey |
There was a problem hiding this comment.
this one need to stay here similar to the DifferentialEvolutionSearch since autotuner/init.py usage results in partial initialization
| args: tuple[object, ...], | ||
| signature: tuple[Hashable, ...], | ||
| ) -> BoundKernelInMemoryCacheKey: | ||
| from ..autotuner.base_cache import BoundKernelInMemoryCacheKey |
stack-info: PR: #336, branch: oulgen/stack/26
Stacked PRs:
[RFC] Implement basic on disk caching