From bb33fda44d8cc889230698539b8df5a7ba114b67 Mon Sep 17 00:00:00 2001 From: Wing Lian Date: Wed, 22 Oct 2025 21:24:52 -0700 Subject: [PATCH] install flash attention in 2.9.0 base images (#3224) --- docker/Dockerfile-base | 6 ++++-- docker/Dockerfile-uv-base | 6 ++++++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base index 87918cc41..cc209f304 100644 --- a/docker/Dockerfile-base +++ b/docker/Dockerfile-base @@ -47,6 +47,8 @@ RUN git lfs install --skip-repo && \ pip3 install -U --no-cache-dir pydantic==1.10.10 && \ pip3 cache purge -RUN if [ "$PYTORCH_VERSION" = "2.6.0" ] && [ "$CUDA" = "124" ] ; then \ - FLASH_ATTENTION_FORCE_BUILD="TRUE" pip3 install --no-build-isolation flash-attn==2.8.0.post2; \ +RUN if [ "$PYTORCH_VERSION" = "2.9.0" ] && [ "$CUDA" = "128" ] ; then \ + wget https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.17/flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \ + pip3 install --no-cache-dir flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \ + rm flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \ fi diff --git a/docker/Dockerfile-uv-base b/docker/Dockerfile-uv-base index eaa49b9e9..2ca272c6e 100644 --- a/docker/Dockerfile-uv-base +++ b/docker/Dockerfile-uv-base @@ -34,3 +34,9 @@ RUN uv pip install packaging setuptools wheel psutil \ && uv pip install --no-build-isolation "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" \ && uv pip install "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" \ && uv pip install awscli pydantic + +RUN if [ "$PYTORCH_VERSION" = "2.9.0" ] && [ "$CUDA" = "128" ] ; then \ + wget https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.4.17/flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \ + uv pip install --no-cache-dir flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \ + rm flash_attn-2.8.3+cu128torch2.9-cp311-cp311-linux_x86_64.whl; \ + fi