diff --git a/.github/workflows/base.yml b/.github/workflows/base.yml index b2681bb5d..87d6772dd 100644 --- a/.github/workflows/base.yml +++ b/.github/workflows/base.yml @@ -60,6 +60,13 @@ jobs: pytorch: 2.9.0 torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX" dockerfile: "Dockerfile-base" + - cuda: "130" + cuda_version: 13.0.0 + cudnn_version: "" + python_version: "3.11" + pytorch: 2.9.0 + torch_cuda_arch_list: "9.0+PTX" + dockerfile: "Dockerfile-base" # - cuda: "128" # cuda_version: 12.8.1 # cudnn_version: "" @@ -143,6 +150,13 @@ jobs: pytorch: 2.9.0 torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX" dockerfile: "Dockerfile-uv-base" + - cuda: "130" + cuda_version: 13.0.0 + cudnn_version: "" + python_version: "3.11" + pytorch: 2.9.0 + torch_cuda_arch_list: "9.0+PTX" + dockerfile: "Dockerfile-uv-base" steps: - name: Checkout uses: actions/checkout@v4 diff --git a/docker/Dockerfile-base b/docker/Dockerfile-base index cc209f304..a08b5cd4f 100644 --- a/docker/Dockerfile-base +++ b/docker/Dockerfile-base @@ -37,10 +37,14 @@ WORKDIR /workspace RUN python3 -m pip install --upgrade pip && pip3 install -U packaging==23.2 setuptools==75.8.0 wheel && \ python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} torchvision --extra-index-url https://download.pytorch.org/whl/cu$CUDA && \ - CAUSAL_CONV1D_FORCE_CXX11_ABI=TRUE CAUSAL_CONV1D_FORCE_BUILD=TRUE python3 -m pip install --no-cache-dir causal_conv1d==1.5.2 && \ - python3 -m pip install --no-cache-dir "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" && \ python3 -m pip cache purge +RUN if [ "$CUDA" != "130" ] ; then \ + CAUSAL_CONV1D_FORCE_CXX11_ABI=TRUE CAUSAL_CONV1D_FORCE_BUILD=TRUE python3 -m pip install --no-cache-dir "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@v1.5.4"; \ + python3 -m pip install --no-cache-dir "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main"; \ + python3 -m pip cache purge; \ + fi + RUN git lfs install --skip-repo && \ pip3 install awscli && \ # The base image ships with `pydantic==1.8.2` which is not working diff --git a/setup.py b/setup.py index 2845bb151..b16377e92 100644 --- a/setup.py +++ b/setup.py @@ -162,7 +162,7 @@ extras_require = { "llmcompressor": [ "llmcompressor==0.5.1", ], - "fbgemm-gpu": ["fbgemm-gpu-genai>=1.2.0"], + "fbgemm-gpu": ["fbgemm-gpu-genai==1.3.0"], "opentelemetry": [ "opentelemetry-api", "opentelemetry-sdk",