From cdc12c307dbedd2182f3e304e2b0e9774078c44a Mon Sep 17 00:00:00 2001 From: Lequn Chen Date: Tue, 29 Oct 2024 13:15:15 -0700 Subject: [PATCH] Change workspace dir (#566) In my development setup, I have different machines with different GPUs. They share the home directory on a network filesystem. When I switch between machines, since the JIT compilation flags change, I'll have to recompile kernels every time. One solution is to specify the same `TORCH_CUDA_ARCH_LIST` every time. However, I keep forgetting that. Another solution, as proposed in this PR, is to put different arch list in different cache directory. --- python/flashinfer/jit/env.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/python/flashinfer/jit/env.py b/python/flashinfer/jit/env.py index e3fbec81..cc48182c 100644 --- a/python/flashinfer/jit/env.py +++ b/python/flashinfer/jit/env.py @@ -15,9 +15,20 @@ """ import pathlib +import re + +from torch.utils.cpp_extension import _get_cuda_arch_flags + + +def _get_workspace_dir_name() -> pathlib.Path: + flags = _get_cuda_arch_flags() + arch = "_".join(sorted(set(re.findall(r"compute_(\d+)", "".join(flags))))) + # e.g.: $HOME/.cache/flashinfer/75_80_89_90/ + return pathlib.Path.home() / ".cache" / "flashinfer" / arch + # use pathlib -FLASHINFER_WORKSPACE_DIR = pathlib.Path.home() / ".flashinfer" +FLASHINFER_WORKSPACE_DIR = _get_workspace_dir_name() FLASHINFER_JIT_DIR = FLASHINFER_WORKSPACE_DIR / "cached_ops" FLASHINFER_GEN_SRC_DIR = FLASHINFER_WORKSPACE_DIR / "generated" _project_root = pathlib.Path(__file__).resolve().parent.parent.parent