root_scripts/fix_abi.sh

#!/bin/bash
set -x

# Fully clean
pip uninstall flash-attn -y 2>&1 | tail -3
rm -f /opt/conda/lib/python3.11/site-packages/flash_attn_2_cuda*.so
rm -rf /opt/conda/lib/python3.11/site-packages/flash_attn*

# Set environment to force correct ABI
export FLASH_ATTENTION_FORCE_CXX11_ABI=FALSE

# Install from PyPI, letting setup.py handle it
pip install flash-attn==2.8.3 --no-build-isolation --force-reinstall \
    2>&1 | tail -20

# Check
python3 -c "from flash_attn import flash_attn_func; print('flash_attn OK')" 2>&1

echo "FIX_DONE"