fix: handle None from cupy.cuda.get_cuda_path() in cuda_launch

cupy.cuda.get_cuda_path() can return None when CUDA_HOME is not set
and cupy can't auto-detect it. Fall back to /usr/local/cuda.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-04-11 02:20:20 +02:00
parent c08fe58fe7
commit 2e75e2d076
2 changed files with 15 additions and 6 deletions
@@ -268,11 +268,14 @@ _cuda_launch_cache = {}
@torch.compiler.disable() @torch.compiler.disable()
def cuda_launch(strKey: str): def cuda_launch(strKey: str):
if strKey not in _cuda_launch_cache: if strKey not in _cuda_launch_cache:
try: if "CUDA_HOME" not in os.environ:
os.environ.setdefault("CUDA_HOME", cupy.cuda.get_cuda_path()) try:
except Exception: cuda_path = cupy.cuda.get_cuda_path()
if "CUDA_HOME" not in os.environ: except Exception:
raise RuntimeError("'CUDA_HOME' not set, unable to find cuda-toolkit installation.") cuda_path = None
if cuda_path is None:
cuda_path = "/usr/local/cuda"
os.environ["CUDA_HOME"] = cuda_path
strKernel = objCudacache[strKey]["strKernel"] strKernel = objCudacache[strKey]["strKernel"]
strFunction = objCudacache[strKey]["strFunction"] strFunction = objCudacache[strKey]["strFunction"]
_cuda_launch_cache[strKey] = cupy.RawModule( _cuda_launch_cache[strKey] = cupy.RawModule(
+7 -1
View File
@@ -224,7 +224,13 @@ _cuda_launch_cache = {}
def cuda_launch(strKey:str): def cuda_launch(strKey:str):
if strKey not in _cuda_launch_cache: if strKey not in _cuda_launch_cache:
if 'CUDA_HOME' not in os.environ: if 'CUDA_HOME' not in os.environ:
os.environ['CUDA_HOME'] = cupy.cuda.get_cuda_path() try:
cuda_path = cupy.cuda.get_cuda_path()
except Exception:
cuda_path = None
if cuda_path is None:
cuda_path = '/usr/local/cuda'
os.environ['CUDA_HOME'] = cuda_path
_cuda_launch_cache[strKey] = cupy.RawKernel( _cuda_launch_cache[strKey] = cupy.RawKernel(
objCudacache[strKey]['strKernel'], objCudacache[strKey]['strKernel'],
objCudacache[strKey]['strFunction'], objCudacache[strKey]['strFunction'],