From ab4b863546438fe3831f9aa76a7b9d42768ba16a Mon Sep 17 00:00:00 2001
From: Rain Jiang <96632942+rainj-me@users.noreply.github.com>
Date: Wed, 11 Mar 2026 13:37:33 -0700
Subject: [PATCH] =?UTF-8?q?fix=20ci=20by=20removing=20nvidia-cutlass-dsl-l?=
 =?UTF-8?q?ibs-base=20and=20force=20reinstall=20n=E2=80=A6=20(#20380)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 docker/Dockerfile                        | 5 +++++
 scripts/ci/cuda/ci_install_dependency.sh | 2 ++
 2 files changed, 7 insertions(+)

diff --git a/docker/Dockerfile b/docker/Dockerfile
index a0a837ef4..055535f85 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -321,6 +321,11 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     && rm -rf /root/.cargo /root/.rustup target dist ~/.cargo \
     && sed -i '/\.cargo\/env/d' /root/.profile /root/.bashrc 2>/dev/null || true
 
+# Pin nvidia-cutlass-dsl to 4.3.5
+RUN --mount=type=cache,target=/root/.cache/pip \
+   python3 -m pip uninstall nvidia-cutlass-dsl-libs-base; \
+   python3 -m pip install nvidia-cutlass-dsl==4.3.5 --force-reinstall;
+
 # Patching packages for CUDA 12/13 compatibility
 # TODO: Remove when torch version covers these packages
 RUN --mount=type=cache,target=/root/.cache/pip if [ "${CUDA_VERSION%%.*}" = "12" ]; then \
diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh
index c6fa958e9..7422bb765 100755
--- a/scripts/ci/cuda/ci_install_dependency.sh
+++ b/scripts/ci/cuda/ci_install_dependency.sh
@@ -330,10 +330,12 @@ bash "${SCRIPT_DIR}/ci_download_flashinfer_cubin.sh"
 
 # Clean nvidia-cutlass-dsl-libs-base for cutedsl lower than 0.4.4
 $PIP_UNINSTALL_CMD nvidia-cutlass-dsl-libs-base $PIP_UNINSTALL_SUFFIX || true
+$PIP_CMD install nvidia-cutlass-dsl==4.3.5 --force-reinstall $PIP_INSTALL_SUFFIX || true
 
 # Show current packages
 $PIP_CMD list
 python3 -c "import torch; print(torch.version.cuda)"
+python3 -c "import cutlass; import cutlass.cute;"
 
 # Prepare the CI runner (cleanup HuggingFace cache, etc.)
 bash "${SCRIPT_DIR}/prepare_runner.sh"