scripts/build_sage_attention.py

import os
import subprocess
import sys

def run_command(command, env=None):
    """
    Runs a command with a specified environment, prints its output, 
    and raises an exception on failure.
    """
    print(f"🚀 Running command: {' '.join(command)}")
    result = subprocess.run(
        command,
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        encoding='utf-8',
        errors='replace'
    )

    if result.stdout:
        print("--- Pip Output ---")
        print(result.stdout.strip())
        print("------------------")

    if result.returncode != 0:
        raise subprocess.CalledProcessError(result.returncode, command)

def install_sage_attention():
    """
    Installs the sageattention package from PyPI using pip, ensuring the
    correct CUDA architecture is set for any potential on-the-fly compilation.
    """
    print("--- [SageAttention Install] Starting installation using pip ---")

    build_env = os.environ.copy()
    build_env["TORCH_CUDA_ARCH_LIST"] = "9.0"
    print(f"🔧 Setting build environment variable: TORCH_CUDA_ARCH_LIST='{build_env['TORCH_CUDA_ARCH_LIST']}'")

    install_command = [
        sys.executable,
        "-m",
        "pip",
        "install",
        "sageattention==2.2.0",
        "--no-build-isolation",
    ]

    try:
        run_command(install_command, env=build_env)
        print("🎉 SageAttention installed successfully via pip!")
    except subprocess.CalledProcessError as e:
        print(f"❌ SageAttention installation via pip failed (Exit Code: {e.returncode}).")
        print("   The application will continue with the default attention mechanism.")
        raise e
    except Exception as e:
        print(f"❌ An unexpected error occurred during pip installation: {e}")
        raise e


if __name__ == "__main__":
    try:
        install_sage_attention()
    except Exception:
        print("\nInstallation script finished with an error.")
        sys.exit(1)