Spaces:
Running on Zero
Running on Zero
File size: 1,371 Bytes
b701455 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | #!/usr/bin/env python3
"""
Patch for SageAttention setup.py to support TORCH_CUDA_ARCH_LIST environment variable.
This allows building without GPUs present during build time.
"""
import sys
setup_py_path = "setup.py"
# Read the original setup.py
with open(setup_py_path, 'r') as f:
content = f.read()
# Find the line where compute_capabilities is initialized
target_line = "compute_capabilities = set()"
if target_line not in content:
print("ERROR: Could not find target line in setup.py")
sys.exit(1)
# Add our patch right after compute_capabilities initialization
patch_code = '''
# Check for TORCH_CUDA_ARCH_LIST environment variable first (Docker build support)
env_arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST", None)
if env_arch_list:
print(f"Using TORCH_CUDA_ARCH_LIST from environment: {env_arch_list}")
arch_list = env_arch_list.replace(" ", ";").split(";")
for arch in arch_list:
arch = arch.strip()
if not arch:
continue
if arch.endswith("+PTX"):
arch = arch[:-4].strip()
if arch:
compute_capabilities.add(arch)
'''
# Insert the patch
content = content.replace(
target_line,
target_line + patch_code
)
# Write back
with open(setup_py_path, 'w') as f:
f.write(content)
print("✓ Successfully patched setup.py to support TORCH_CUDA_ARCH_LIST")
|