开发者
资源
我要评分
获取效率
正确性
完整性
易理解
在线提单
论坛求助

源码编译构建

获取FlashAttention源码,从源码构建wheel并安装本地构建产物。

  1. 获取FlashAttention 2.8.3源码。
    1
    git clone --branch v2.8.3 --depth 1 https://github.com/Dao-AILab/flash-attention.git
    
  2. 确认PyTorch为CUDA版本。
     1
     2
     3
     4
     5
     6
     7
     8
     9
    10
    11
    cd flash-attention
    python3 - <<'PY'
    import torch  
    
    print("torch_version=" + torch.__version__) 
    print("torch_cuda=" + str(torch.version.cuda)) 
    print("cuda_built=" + str(torch.backends.cuda.is_built())) 
    print("cuda_available=" + str(torch.cuda.is_available())) 
    assert torch.backends.cuda.is_built() 
    assert torch.cuda.is_available() 
    PY
    
  3. 从源码构建wheel。
    1
    2
    3
    4
    5
    6
    7
    8
    9
    export CUDA_HOME=/usr/local/cuda-13.0 
    export PATH=$CUDA_HOME/bin:$PATH
    export LD_LIBRARY_PATH=$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}
    export FLASH_ATTENTION_FORCE_BUILD=TRUE 
    export FLASH_ATTN_CUDA_ARCHS=80 
    export MAX_JOBS=8 
    export NVCC_THREADS=1  
    
    python3 -m pip wheel --no-build-isolation --no-deps --wheel-dir dist .
    
  4. 安装本地构建产物。
    1
    python3 -m pip install --force-reinstall --no-deps dist/flash_attn-2.8.3-*.whl
    
    构建产物示例。
    1
    flash_attn-2.8.3-cp311-cp311-linux_aarch64.whl