Check for _is_cuda() in compute_num_jobs (vllm-project#3481)

Temirulan · Mar 20, 2024 · 9f5b645 · 9f5b645
1 parent fd39980
commit 9f5b645
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/setup.py b/setup.py
@@ -61,12 +61,12 @@ def compute_num_jobs(self):
         except AttributeError:
             num_jobs = os.cpu_count()
 
-        nvcc_cuda_version = get_nvcc_cuda_version()
-        if nvcc_cuda_version >= Version("11.2"):
-            nvcc_threads = int(os.getenv("NVCC_THREADS", 8))
-            num_jobs = max(1, round(num_jobs / (nvcc_threads / 4)))
-        else:
-            nvcc_threads = None
+        nvcc_threads = None
+        if _is_cuda():
+            nvcc_cuda_version = get_nvcc_cuda_version()
+            if nvcc_cuda_version >= Version("11.2"):
+                nvcc_threads = int(os.getenv("NVCC_THREADS", 8))
+                num_jobs = max(1, round(num_jobs / (nvcc_threads / 4)))
 
         return num_jobs, nvcc_threads