Skip to content

Fix flash-attention build env in CI deps Dockerfile #565

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
VeeraRajasekhar merged 3 commits into from
Apr 29, 2026
+11 −6
Merged

Fix flash-attention build env in CI deps Dockerfile #565

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a4c0942
Fix flash-attention build in CI deps Dockerfile; bump default TE CI i…
VeeraRajasekhar Apr 27, 2026
a30316e
Addressed reviews
VeeraRajasekhar Apr 28, 2026
ba241bf
Addressed reviews
VeeraRajasekhar Apr 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 11 additions & 6 deletions .github/scripts/Dockerfile.ci.deps
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,13 @@ ARG BASE_DOCKER=registry-sc-harbor.amd.com/framework/compute-rocm-rel-7.2:57_ubu
FROM $BASE_DOCKER
WORKDIR /

# Updated git via git-core PPA
RUN apt-get update && apt-get install -y --no-install-recommends software-properties-common \
&& add-apt-repository ppa:git-core/ppa -y \
&& apt-get update \
&& apt-get install -y --no-install-recommends git vim \
&& rm -rf /var/lib/apt/lists/*

# Build arguments
ARG FA_VERSION=v2.8.1
ARG ROCM_VERSION=7.2
Expand All @@ -17,21 +24,19 @@ ARG AITER_COMMIT=77455e3ecf4f0d28756afc452e914940c45b944b

RUN pip install setuptools wheel
RUN pip install ipython pytest fire pydantic pybind11 ninja pandas
RUN apt-get update && apt-get install -y vim

# Install flash-attention
ENV GPU_ARCHS=gfx950;gfx942
RUN git clone --branch ${FA_VERSION} --depth 1 https://github.com/Dao-AILab/flash-attention.git \
&& cd flash-attention \
&& FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE && FLASH_ATTENTION_SKIP_CK_BUILD=FALSE python setup.py install \
&& GPU_ARCHS="gfx950;gfx942" FLASH_ATTENTION_TRITON_AMD_ENABLE=TRUE FLASH_ATTENTION_SKIP_CK_BUILD=FALSE python setup.py install \
&& cd ..

# AITER (submodule: composable_kernel / CK). Submodule update after checkout aligns CK with this commit.
RUN git clone --recursive https://github.com/ROCm/aiter.git \
# Install AITER
RUN git clone --no-checkout https://github.com/ROCm/aiter.git \
&& cd aiter \
&& git checkout ${AITER_COMMIT} \
&& git submodule update --init --recursive \
&& PIP_FIND_LINKS="https://rocm.frameworks-nightlies.amd.com/whl/gfx942-gfx950/flydsl/" PIP_PRE=1 GPU_ARCHS=native pip install .
&& pip install .

# Install JAX
RUN ROCM_MAJOR=$(echo "${ROCM_VERSION}" | cut -d. -f1) && pip install \
Expand Down
Loading