Manually specify flags if no arch set

stack-info: PR: #2219, branch: drisspg/stack/55

Author: drisspg

CUDA_ARCH_NOTES.md

@@ -79,6 +79,35 @@ def use_debug_mode():
     _get_cuda_arch_flags,
 )
 
+# =====================================================================================
+# CUDA Architecture Settings
+# =====================================================================================
+# If TORCH_CUDA_ARCH_LIST is not set during compilation, PyTorch tries to automatically
+# detect architectures from available GPUs. This can fail when:
+#   1. No GPU is visible to PyTorch
+#   2. CUDA is available but no device is detected
+#
+# To resolve this, you can manually set CUDA architecture targets:
+#   export TORCH_CUDA_ARCH_LIST="7.0;7.5;8.0;8.6+PTX"
+#
+# Adding "+PTX" to the last architecture enables JIT compilation for future GPUs.
+# =====================================================================================
+if "TORCH_CUDA_ARCH_LIST" not in os.environ and torch.version.cuda:
+    # Set to common architectures for CUDA 12.x compatibility
+    cuda_arch_list = "7.0;7.5;8.0;8.6;8.9;9.0"
+    
+    # Only add SM10.0 (Blackwell) flags when using CUDA 12.8 or newer
+    cuda_version = torch.version.cuda
+    if cuda_version and cuda_version.startswith("12.8"):
+        print("Detected CUDA 12.8 - adding SM10.0 architectures to build list")
+        cuda_arch_list += ";10.0"
+    
+    # Add PTX to the last architecture for future compatibility
+    cuda_arch_list += "+PTX"
+    
+    os.environ["TORCH_CUDA_ARCH_LIST"] = cuda_arch_list
+    print(f"Setting default TORCH_CUDA_ARCH_LIST={os.environ['TORCH_CUDA_ARCH_LIST']}")
+
 IS_ROCM = (torch.version.hip is not None) and (ROCM_HOME is not None)