bump fa4 to official released fa4 pkg (#20303)

This commit is contained in:
Rain Jiang
2026-03-17 17:22:56 -07:00
committed by GitHub
parent c77d7c629e
commit cb1e63aba4
4 changed files with 11 additions and 25 deletions

View File

@@ -321,22 +321,21 @@ 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;
python3 -m pip install "nvidia-cutlass-dsl>=4.4.1" "nvidia-cutlass-dsl-libs-base>=4.4.1" --force-reinstall --no-deps;
# 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 \
python3 -m pip install nvidia-nccl-cu12==2.28.3 --force-reinstall --no-deps ; \
python3 -m pip install nvidia-cudnn-cu12==9.16.0.29 --force-reinstall --no-deps ; \
python3 -m pip install cuda-python==12.9 ; \
elif [ "${CUDA_VERSION%%.*}" = "13" ]; then \
python3 -m pip install nvidia-nccl-cu13==2.28.3 --force-reinstall --no-deps ; \
python3 -m pip install nvidia-cudnn-cu13==9.16.0.29 --force-reinstall --no-deps ; \
python3 -m pip install nvidia-cublas==13.1.0.3 --force-reinstall --no-deps ; \
python3 -m pip install nixl-cu13 --no-deps ; \
python3 -m pip install cuda-python==13.1.1 ; \
python3 -m pip install cuda-python==13.2.0 ; \
fi
# Install development tools