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:
@@ -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:
|
|
||||||
os.environ.setdefault("CUDA_HOME", cupy.cuda.get_cuda_path())
|
|
||||||
except Exception:
|
|
||||||
if "CUDA_HOME" not in os.environ:
|
if "CUDA_HOME" not in os.environ:
|
||||||
raise RuntimeError("'CUDA_HOME' not set, unable to find cuda-toolkit installation.")
|
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
|
||||||
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(
|
||||||
|
|||||||
@@ -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'],
|
||||||
|
|||||||
Reference in New Issue
Block a user