diff --git a/pkgs/development/python-modules/deep-gemm/default.nix b/pkgs/development/python-modules/deep-gemm/default.nix index 6c873f09ea20..4946efaf98ef 100644 --- a/pkgs/development/python-modules/deep-gemm/default.nix +++ b/pkgs/development/python-modules/deep-gemm/default.nix @@ -2,6 +2,8 @@ lib, buildPythonPackage, fetchFromGitHub, + replaceVars, + symlinkJoin, # build-system setuptools, @@ -47,6 +49,23 @@ buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: { patches = [ ./use-system-libraries.patch + + # DeepGEMM does JIT compilation and needs to access the NVIDIA compiler and some libraries at + # runtime. + # Instead of letting it search for the cuda toolkit on the host, hardcode the path to a custom + # closure. + (replaceVars ./patch-runtime-cuda-home-path.patch { + cuda_home = symlinkJoin { + name = "cuda-toolkit"; + paths = with cudaPackages; [ + (lib.getBin cuda_nvcc) # bin/nvcc, bin/ptxas, nvvm/, nvcc.profile + (lib.getBin cutlass) # include/cute, include/cutlass + (lib.getInclude cuda_cccl) # include/cuda/std/* (libcu++) + (lib.getInclude cuda_cudart) # include/cuda_runtime.h, cuda_bf16.h, cuda_fp8.h + (lib.getInclude cuda_cuobjdump) # bin/cuobjdump + ]; + }; + }) ]; env = optionalAttrs cudaSupport { diff --git a/pkgs/development/python-modules/deep-gemm/patch-runtime-cuda-home-path.patch b/pkgs/development/python-modules/deep-gemm/patch-runtime-cuda-home-path.patch new file mode 100644 index 000000000000..6391bdf0fae7 --- /dev/null +++ b/pkgs/development/python-modules/deep-gemm/patch-runtime-cuda-home-path.patch @@ -0,0 +1,27 @@ +diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py +index a4633ae..6508865 100644 +--- a/deep_gemm/__init__.py ++++ b/deep_gemm/__init__.py +@@ -62,21 +62,7 @@ from .utils import * + + # Initialize CPP modules + def _find_cuda_home() -> str: +- # TODO: reuse PyTorch API later +- # For some PyTorch versions, the original `_find_cuda_home` will initialize CUDA, which is incompatible with process forks +- cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH') +- if cuda_home is None: +- # noinspection PyBroadException +- try: +- with open(os.devnull, 'w') as devnull: +- nvcc = subprocess.check_output(['which', 'nvcc'], stderr=devnull).decode().rstrip('\r\n') +- cuda_home = os.path.dirname(os.path.dirname(nvcc)) +- except Exception: +- cuda_home = '/usr/local/cuda' +- if not os.path.exists(cuda_home): +- cuda_home = None +- assert cuda_home is not None +- return cuda_home ++ return "@cuda_home@" + + + deep_gemm_cpp.init(