diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 934a6b996..3b8dd0bbb 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -22,6 +22,7 @@ ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" ENV AITER_COMMIT="v0.1.4" ENV NO_DEPS_FLAG="" +ENV AITER_MXFP4_MOE_SF="0" # =============================== # Base image 942 and args @@ -33,6 +34,7 @@ ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" ENV AITER_COMMIT="v0.1.7.post1" ENV NO_DEPS_FLAG="" +ENV AITER_MXFP4_MOE_SF="0" # =============================== # Base image 950 and args @@ -44,7 +46,7 @@ ENV BUILD_AITER_ALL="1" ENV BUILD_MOONCAKE="1" ENV AITER_COMMIT="v0.1.7.post1" ENV NO_DEPS_FLAG="" - +ENV AITER_MXFP4_MOE_SF="1" # =============================== # Chosen arch and args FROM ${GPU_ARCH} @@ -106,6 +108,7 @@ RUN git clone ${AITER_REPO} \ && git submodule update --init --recursive RUN cd aiter \ && if [ "$GPU_ARCH" = "gfx950" ]; then export AITER_MXFP4_MOE_SF=1; fi \ + && echo "[AITER] GPU_ARCH=${GPU_ARCH} AITER_MXFP4_MOE_SF=${AITER_MXFP4_MOE_SF:-unset}" \ && if [ "$BUILD_AITER_ALL" = "1" ] && [ "$BUILD_LLVM" = "1" ]; then \ sh -c "HIP_CLANG_PATH=/sgl-workspace/llvm-project/build/bin/ PREBUILD_KERNELS=1 GPU_ARCHS=$GPU_ARCH_LIST python setup.py develop"; \ elif [ "$BUILD_AITER_ALL" = "1" ]; then \ @@ -296,12 +299,7 @@ RUN python3 -m pip install --no-cache-dir \ # ----------------------- # Performance environment variable. -# Set AITER_MXFP4_MOE_SF only for gfx950 architecture -RUN if [ "$GPU_ARCH" = "gfx950" ]; then \ - echo "AITER_MXFP4_MOE_SF=1"; \ - else \ - echo "AITER_MXFP4_MOE_SF=0"; \ - fi >> /etc/environment +RUN echo "AITER_MXFP4_MOE_SF=${AITER_MXFP4_MOE_SF}" >> /etc/environment ENV HIP_FORCE_DEV_KERNARG=1 ENV HSA_NO_SCRATCH_RECLAIM=1