root_scripts/install_flash_attn.sh

#!/bin/bash
set -x

echo "=== Trying flash-attn install ==="

# Method 1: Try FA3 from the official wheel index
echo "--- Method 1: FA3 wheel index ---"
pip install flash-attn \
    --extra-index-url https://dao-ailab.github.io/flash-attention/wheels/ \
    --no-build-isolation \
    2>&1 | tail -20

# Check if it worked
python3 -c "import flash_attn; print(f'flash_attn version: {flash_attn.__version__}')" 2>&1

echo "INSTALL_SCRIPT_DONE"