Skip to content

Commit

Permalink
Change workspace dir (#566)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
abcdabcd987 authored Oct 29, 2024
1 parent d30667b commit cdc12c3
Showing 1 changed file with 12 additions and 1 deletion.
13 changes: 12 additions & 1 deletion python/flashinfer/jit/env.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit cdc12c3

Please # to comment.