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"