python3Packages.deep-gemm: patch runtime cudatoolkit discovery logic

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 {

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(