pip install flash-attn hangs on Google Colab A100 with latest environment

Hello,

The command pip install flash-attn --no-build-isolation is currently hanging indefinitely on a standard Google Colab A100 instance. It seems a pre-compiled wheel is not available for the current Colab environment, and building from source fails.

This blocks the use of attn_implementation="flash_attention_2" in libraries like Hugging Face transformers.

Environment Details:

GPU: NVIDIA A100 | NVIDIA-SMI 550.54.15   Driver Version: 550.54.15 

Python Version: Python 3.12.11

CUDA Version: 12.4

PyTorch Version: 2.3.1+cu121

PyTorch CUDA Version: 12.1

Steps to Reproduce:

Start a new Google Colab notebook and select an A100 GPU runtime.

Run the following cell to install PyTorch:
  !pip install torch==2.3.1 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu121

Run the installation for FlashAttention: !pip install flash-attn --no-build-isolation

Observed Behavior: The pip install flash-attn command hangs and never completes.

Expected Behavior: The command should quickly install a pre-compiled wheel compatible with the environment.

Thank you!