flash_attn build for NVidia DGX Spark within a uv venv.

install_flash_attn.sh

#!/bin/bash
# Helper script to install flash-attn with proper CUDA environment setup
# This script handles the CUDA_HOME environment variable and version compatibility

set -e

# Check for Python development headers
echo "Checking for Python development headers..."
PYTHON_VERSION=$(python3 --version | cut -d' ' -f2 | cut -d'.' -f1,2)
PYTHON_INCLUDE="/usr/include/python${PYTHON_VERSION}/Python.h"

# Also check using find as fallback
if [ ! -f "$PYTHON_INCLUDE" ]; then
    PYTHON_H_FOUND=$(find /usr/include -name "Python.h" 2>/dev/null | head -1)
    if [ -n "$PYTHON_H_FOUND" ] && [ -f "$PYTHON_H_FOUND" ]; then
        echo "✓ Python development headers found at: $PYTHON_H_FOUND"
    else
        echo "❌ Error: Python development headers not found!"
        echo "   Expected: $PYTHON_INCLUDE"
        echo "   Please install with: sudo apt-get install python${PYTHON_VERSION}-dev"
        echo ""
        echo "   Or for Python 3.12: sudo apt-get install python3.12-dev"
        exit 1
    fi
else
    echo "✓ Python development headers found at: $PYTHON_INCLUDE"
fi

# Set CUDA environment
export CUDA_HOME=/usr/local/cuda
export PATH=/usr/local/cuda/bin:$PATH

# Verify CUDA is available
if [ ! -f "$CUDA_HOME/bin/nvcc" ]; then
    echo "Error: nvcc not found at $CUDA_HOME/bin/nvcc"
    exit 1
fi

echo "CUDA_HOME: $CUDA_HOME"
echo "nvcc version:"
$CUDA_HOME/bin/nvcc --version | head -n 4

# Activate virtual environment if it exists
if [ -d ".venv" ]; then
    source .venv/bin/activate
fi

# Patch torch cpp_extension to skip CUDA version check (for CUDA 13.0 compatibility)
python3 << 'PYTHON_SCRIPT'
import torch.utils.cpp_extension as cpp_ext
import inspect
import os

file_path = inspect.getfile(cpp_ext._check_cuda_version)
print(f"Checking torch cpp_extension: {file_path}")

with open(file_path, 'r') as f:
    content = f.read()

if '# PATCHED: Skip CUDA version check' not in content:
    # Create backup
    with open(file_path + '.bak', 'w') as f:
        f.write(content)
    
    # Patch the function
    patched_content = content.replace(
        'def _check_cuda_version(compiler_name, compiler_version):',
        'def _check_cuda_version(compiler_name, compiler_version):\n    # PATCHED: Skip CUDA version check for CUDA 13.0 compatibility\n    return'
    )
    
    with open(file_path, 'w') as f:
        f.write(patched_content)
    print("✓ Patched torch cpp_extension to skip CUDA version check")
else:
    print("✓ torch cpp_extension already patched")
PYTHON_SCRIPT

# Ensure CUDA-enabled torch is installed
echo "Checking torch installation..."
python3 << 'PYTHON_SCRIPT'
import torch
print(f"Current torch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if not torch.cuda.is_available() or '+cpu' in torch.__version__:
    print("⚠ Warning: Torch is CPU-only. Reinstalling with CUDA support...")
    import subprocess
    subprocess.check_call([
        "uv", "pip", "install", "--force-reinstall",
        "torch", "torchvision", "torchaudio",
        "--index-url", "https://download.pytorch.org/whl/cu124"
    ])
PYTHON_SCRIPT

# Install flash-attn with CUDA_HOME set
echo "Installing flash-attn..."
echo "Note: This may take 10-30 minutes and requires significant memory/CPU resources."
echo "Setting MAX_JOBS=2 to reduce memory pressure during compilation..."
export MAX_JOBS=2
echo ""
CUDA_HOME=/usr/local/cuda PATH=/usr/local/cuda/bin:$PATH uv pip install flash-attn --no-build-isolation

echo "✓ flash-attn installed successfully!"