From 2ce2377740ad90357a3f8a8a3d79fd42f7737316 Mon Sep 17 00:00:00 2001
From: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Date: Thu, 13 Nov 2025 17:03:41 -0800
Subject: [PATCH] [AMD] Fix AITER_MXFP4_MOE_SF setting for gfx950 (#13239)

---
 docker/rocm.Dockerfile | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

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