From fda9ceadef707d0c1dc1c828e6611127f86be310 Mon Sep 17 00:00:00 2001 From: Alexander Grund Date: Wed, 11 Mar 2026 17:29:30 +0100 Subject: [PATCH] Don't detect local GPU if `$DS_IGNORE_CUDA_DETECTION` is set Make this consistent over all OPs: For cross-compilation we should not check the local GPU version. --- op_builder/fp_quantizer.py | 24 +++++++++++++----------- op_builder/inference_core_ops.py | 23 ++++++++++++----------- op_builder/inference_cutlass_builder.py | 23 ++++++++++++----------- op_builder/ragged_ops.py | 23 ++++++++++++----------- op_builder/ragged_utils.py | 23 ++++++++++++----------- op_builder/spatial_inference.py | 20 +++++++++++--------- op_builder/transformer_inference.py | 24 +++++++++++++----------- 7 files changed, 85 insertions(+), 75 deletions(-) diff --git a/op_builder/fp_quantizer.py b/op_builder/fp_quantizer.py index 5ccc35ac2b1f..d6e0a51b1ece 100644 --- a/op_builder/fp_quantizer.py +++ b/op_builder/fp_quantizer.py @@ -3,6 +3,7 @@ # DeepSpeed Team +import os try: from packaging import version as pkg_version except ImportError: @@ -31,19 +32,20 @@ def is_compatible(self, verbose=False): return False cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda - if cuda_capability < 8: - if verbose: - self.warning("NVIDIA Inference is only supported on Ampere and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: + if not os.environ.get("DS_IGNORE_CUDA_DETECTION"): + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 8: if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") + self.warning("NVIDIA Inference is only supported on Ampere and newer architectures") cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False try: import triton diff --git a/op_builder/inference_core_ops.py b/op_builder/inference_core_ops.py index b6665ebb7618..d3b0f3aaeeb9 100755 --- a/op_builder/inference_core_ops.py +++ b/op_builder/inference_core_ops.py @@ -28,19 +28,20 @@ def is_compatible(self, verbose=False): return False cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda - if cuda_capability < 6: - if verbose: - self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: + if not os.environ.get("DS_IGNORE_CUDA_DETECTION"): + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False return super().is_compatible(verbose) and cuda_okay def filter_ccs(self, ccs): diff --git a/op_builder/inference_cutlass_builder.py b/op_builder/inference_cutlass_builder.py index a4a607288ca8..5b2299e9c5cc 100644 --- a/op_builder/inference_cutlass_builder.py +++ b/op_builder/inference_cutlass_builder.py @@ -27,19 +27,20 @@ def is_compatible(self, verbose=False): return False cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda - if cuda_capability < 6: - if verbose: - self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: + if not os.environ.get("DS_IGNORE_CUDA_DETECTION"): + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False return super().is_compatible(verbose) and cuda_okay def filter_ccs(self, ccs): diff --git a/op_builder/ragged_ops.py b/op_builder/ragged_ops.py index 0df28cc2282a..d96f437e16ca 100644 --- a/op_builder/ragged_ops.py +++ b/op_builder/ragged_ops.py @@ -28,19 +28,20 @@ def is_compatible(self, verbose=False): return False cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda - if cuda_capability < 6: - if verbose: - self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: + if not os.environ.get("DS_IGNORE_CUDA_DETECTION"): + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False return super().is_compatible(verbose) and cuda_okay def filter_ccs(self, ccs): diff --git a/op_builder/ragged_utils.py b/op_builder/ragged_utils.py index 208c9f833ebe..6cf4e4b3153d 100755 --- a/op_builder/ragged_utils.py +++ b/op_builder/ragged_utils.py @@ -28,19 +28,20 @@ def is_compatible(self, verbose=False): return False cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda - if cuda_capability < 6: - if verbose: - self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: + if not os.environ.get("DS_IGNORE_CUDA_DETECTION"): + if not self.is_rocm_pytorch() and torch.cuda.is_available(): #ignore-cuda + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major #ignore-cuda + if cuda_capability < 6: if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False return super().is_compatible(verbose) and cuda_okay def filter_ccs(self, ccs): diff --git a/op_builder/spatial_inference.py b/op_builder/spatial_inference.py index d6c5fa661156..57714c8c6bf5 100644 --- a/op_builder/spatial_inference.py +++ b/op_builder/spatial_inference.py @@ -3,6 +3,7 @@ # DeepSpeed Team +import os from .builder import CUDAOpBuilder, installed_cuda_version @@ -26,15 +27,16 @@ def is_compatible(self, verbose=False): return False cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: - if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") - cuda_okay = False + if not os.environ.get("DS_IGNORE_CUDA_DETECTION"): + if not self.is_rocm_pytorch() and torch.cuda.is_available(): + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False return super().is_compatible(verbose) and cuda_okay def sources(self): diff --git a/op_builder/transformer_inference.py b/op_builder/transformer_inference.py index 3afa74dc31c2..2507ee4ee692 100755 --- a/op_builder/transformer_inference.py +++ b/op_builder/transformer_inference.py @@ -3,6 +3,7 @@ # DeepSpeed Team +import os from .builder import CUDAOpBuilder, installed_cuda_version @@ -26,19 +27,20 @@ def is_compatible(self, verbose=False): return False cuda_okay = True - if not self.is_rocm_pytorch() and torch.cuda.is_available(): - sys_cuda_major, _ = installed_cuda_version() - torch_cuda_major = int(torch.version.cuda.split('.')[0]) - cuda_capability = torch.cuda.get_device_properties(0).major - if cuda_capability < 6: - if verbose: - self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") - cuda_okay = False - if cuda_capability >= 8: - if torch_cuda_major < 11 or sys_cuda_major < 11: + if not os.environ.get("DS_IGNORE_CUDA_DETECTION"): + if not self.is_rocm_pytorch() and torch.cuda.is_available(): + sys_cuda_major, _ = installed_cuda_version() + torch_cuda_major = int(torch.version.cuda.split('.')[0]) + cuda_capability = torch.cuda.get_device_properties(0).major + if cuda_capability < 6: if verbose: - self.warning("On Ampere and higher architectures please use CUDA 11+") + self.warning("NVIDIA Inference is only supported on Pascal and newer architectures") cuda_okay = False + if cuda_capability >= 8: + if torch_cuda_major < 11 or sys_cuda_major < 11: + if verbose: + self.warning("On Ampere and higher architectures please use CUDA 11+") + cuda_okay = False return super().is_compatible(verbose) and cuda_okay def filter_ccs(self, ccs):