A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[ROCm][Build][Bugfix] Fix ROCm base docker whls installation order (#25415)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Sep 24, 2025 at 11:25 UTC
302eb941f360aaac6936abc782e235b45da88699
1 file changed
+12
-12
docker/Dockerfile.rocm_base
+12
-12
@@ -65,8 +65,6 @@ ARG PYTORCH_BRANCH
65
ARG PYTORCH_VISION_BRANCH
66
ARG PYTORCH_REPO
67
ARG PYTORCH_VISION_REPO
68
-ARG FA_BRANCH
69
-ARG FA_REPO
68
RUN git clone ${PYTORCH_REPO} pytorch
69
RUN cd pytorch && git checkout ${PYTORCH_BRANCH} && \
70
pip install -r requirements.txt && git submodule update --init --recursive \
@@ -77,14 +75,20 @@ RUN git clone ${PYTORCH_VISION_REPO} vision
75
RUN cd vision && git checkout ${PYTORCH_VISION_BRANCH} \
76
&& python3 setup.py bdist_wheel --dist-dir=dist \
77
&& pip install dist/*.whl
78
+RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \
79
+ && cp /app/vision/dist/*.whl /app/install
80
+
81
+FROM base AS build_fa
82
+ARG FA_BRANCH
83
+ARG FA_REPO
84
+RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
85
+ pip install /install/*.whl
86
RUN git clone ${FA_REPO}
87
RUN cd flash-attention \
88
&& git checkout ${FA_BRANCH} \
89
&& git submodule update --init \
90
&& GPU_ARCHS=$(echo ${PYTORCH_ROCM_ARCH} | sed -e 's/;gfx1[0-9]\{3\}//g') python3 setup.py bdist_wheel --dist-dir=dist
85
-RUN mkdir -p /app/install && cp /app/pytorch/dist/*.whl /app/install \
86
- && cp /app/vision/dist/*.whl /app/install \
87
- && cp /app/flash-attention/dist/*.whl /app/install
91
+RUN mkdir -p /app/install && cp /app/flash-attention/dist/*.whl /app/install
92
93
FROM base AS build_aiter
94
ARG AITER_BRANCH
@@ -103,6 +107,8 @@ FROM base AS debs
107
RUN mkdir /app/debs
108
RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
109
cp /install/*.whl /app/debs
110
+RUN --mount=type=bind,from=build_fa,src=/app/install/,target=/install \
111
+ cp /install/*.whl /app/debs
112
RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
113
cp /install/*.whl /app/debs
114
RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
@@ -111,13 +117,7 @@ RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
117
cp /install/*.whl /app/debs
118
119
FROM base AS final
114
-RUN --mount=type=bind,from=build_triton,src=/app/install/,target=/install \
115
- pip install /install/*.whl
116
-RUN --mount=type=bind,from=build_amdsmi,src=/app/install/,target=/install \
117
- pip install /install/*.whl
118
-RUN --mount=type=bind,from=build_pytorch,src=/app/install/,target=/install \
119
- pip install /install/*.whl
120
-RUN --mount=type=bind,from=build_aiter,src=/app/install/,target=/install \
120
+RUN --mount=type=bind,from=debs,src=/app/debs,target=/install \
121
pip install /install/*.whl
122
123
ARG BASE_IMAGE