From 8b0976736909690e61c3b482a7ab6f5d0c6570ce Mon Sep 17 00:00:00 2001 From: Brian Ryu Date: Mon, 2 Feb 2026 19:07:15 +0000 Subject: [PATCH 1/5] Add cu131 container in CI build --- .devcontainer/cu131/devcontainer.json | 26 +++++++++ .github/workflows/release-ci-docker.yml | 6 +- ci/docker-tags.yml | 1 + docker/Dockerfile.cu131.dev | 73 +++++++++++++++++++++++++ 4 files changed, 104 insertions(+), 2 deletions(-) create mode 100644 .devcontainer/cu131/devcontainer.json create mode 100644 docker/Dockerfile.cu131.dev diff --git a/.devcontainer/cu131/devcontainer.json b/.devcontainer/cu131/devcontainer.json new file mode 100644 index 0000000000..ed8de17409 --- /dev/null +++ b/.devcontainer/cu131/devcontainer.json @@ -0,0 +1,26 @@ +{ + "name": "CUDA Development Container", + "build": { + "dockerfile": "../../docker/Dockerfile.cu131.dev", + "context": "../../" + }, + "runArgs": [ + "--gpus=all" + ], + "customizations": { + "vscode": { + "extensions": [ + "llvm-vs-code-extensions.vscode-clangd", + "ms-python.python", + "nvidia.nsight-vscode-edition", + "xaver.clang-format", + "charliermarsh.ruff", + "matangover.mypy" + ] + } + }, + "mounts": [ + "type=bind,source=${localEnv:HOME}/.ssh,target=/home/devuser/.ssh,readonly" + ], + "remoteUser": "devuser" +} diff --git a/.github/workflows/release-ci-docker.yml b/.github/workflows/release-ci-docker.yml index 3fb36129e1..1b158dc1e1 100644 --- a/.github/workflows/release-ci-docker.yml +++ b/.github/workflows/release-ci-docker.yml @@ -36,7 +36,7 @@ jobs: needs: generate-tag strategy: matrix: - cuda: [cu126, cu128, cu129, cu130] + cuda: [cu126, cu128, cu129, cu130, cu131] arch: [amd64, arm64] steps: - uses: actions/checkout@v4 @@ -74,7 +74,7 @@ jobs: needs: [generate-tag, build] strategy: matrix: - cuda: [cu126, cu128, cu129, cu130] + cuda: [cu126, cu128, cu129, cu130, cu131] steps: - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 @@ -111,6 +111,7 @@ jobs: flashinfer/flashinfer-ci-cu128: ${DATE_SHA} flashinfer/flashinfer-ci-cu129: ${DATE_SHA} flashinfer/flashinfer-ci-cu130: ${DATE_SHA} + flashinfer/flashinfer-ci-cu131: ${DATE_SHA} EOF - name: Create Pull Request @@ -127,6 +128,7 @@ jobs: - flashinfer/flashinfer-ci-cu128:${{ needs.generate-tag.outputs.date_sha }} - flashinfer/flashinfer-ci-cu129:${{ needs.generate-tag.outputs.date_sha }} - flashinfer/flashinfer-ci-cu130:${{ needs.generate-tag.outputs.date_sha }} + - flashinfer/flashinfer-ci-cu131:${{ needs.generate-tag.outputs.date_sha }} Auto-generated by [release-ci-docker workflow](https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}) branch: update-docker-tags-${{ needs.generate-tag.outputs.date_sha }} diff --git a/ci/docker-tags.yml b/ci/docker-tags.yml index a2e92e5f1b..d2fb133fee 100644 --- a/ci/docker-tags.yml +++ b/ci/docker-tags.yml @@ -2,3 +2,4 @@ flashinfer/flashinfer-ci-cu126: 20260131-a52eff1 flashinfer/flashinfer-ci-cu128: 20260131-a52eff1 flashinfer/flashinfer-ci-cu129: 20260131-a52eff1 flashinfer/flashinfer-ci-cu130: 20260131-a52eff1 +flashinfer/flashinfer-ci-cu131: 20260131-a52eff1 diff --git a/docker/Dockerfile.cu131.dev b/docker/Dockerfile.cu131.dev new file mode 100644 index 0000000000..8b7ea81e98 --- /dev/null +++ b/docker/Dockerfile.cu131.dev @@ -0,0 +1,73 @@ +FROM nvidia/cuda:13.1.1-devel-ubuntu24.04 + +ENV DEBIAN_FRONTEND=noninteractive + +# Update package lists and install system dependencies +RUN apt-get update && apt-get install -y \ + curl \ + git \ + wget \ + clang-format \ + clangd-19 \ + vim \ + zsh \ + && rm -rf /var/lib/apt/lists/* + +# Create a non-root user +ARG USERNAME=devuser +ARG USER_UID=1003 +ARG USER_GID=$USER_UID + +# Create the user +RUN groupadd --gid $USER_GID $USERNAME \ + && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME \ + # [Optional] Add sudo support + && apt-get update \ + && apt-get install -y sudo \ + && echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \ + && chmod 0440 /etc/sudoers.d/$USERNAME \ + && rm -rf /var/lib/apt/lists/* + +# Remove default 'ubuntu' user (UID 1000) to prevent devcontainer permission conflicts +# Ref: https://github.com/rapidsai/devcontainers/pull/373 +RUN if grep ubuntu:x:1000:1000 /etc/passwd >/dev/null; then userdel -f -r ubuntu; fi + +# Switch to non-root user +USER $USERNAME +WORKDIR /home/$USERNAME + +# Install python +COPY docker/install/install_python.sh /install/install_python.sh +RUN bash /install/install_python.sh /home/$USERNAME/conda py312 + +RUN echo "source activate py312" >> ~/.bashrc +ENV PATH="/home/$USERNAME/conda/bin:$PATH" +ENV PATH="/home/$USERNAME/conda/envs/py312/bin:$PATH" + +# Install torch and other python packages +COPY requirements.txt /install/requirements.txt +COPY docker/install/install_python_packages.sh /install/install_python_packages.sh +RUN bash /install/install_python_packages.sh cu131 && pip3 install pre-commit + +# Install mpi4py in the conda environment +RUN conda install -n py312 -y mpi4py mpich + +# Install oh-my-zsh +RUN sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended + +# Install zsh-autosuggestions +RUN git clone https://github.com/zsh-users/zsh-autosuggestions ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-autosuggestions + +# Configure zsh +RUN sed -i 's/ZSH_THEME="robbyrussell"/ZSH_THEME="fino-time"/' ~/.zshrc && \ + sed -i 's/plugins=(git)/plugins=(git zsh-autosuggestions)/' ~/.zshrc + +# clangd +ENV PATH="/usr/lib/llvm-19/bin:$PATH" + +# Triton +ENV TRITON_PTXAS_PATH="/usr/local/cuda/bin/ptxas" + +# Set zsh as default shell +ENV SHELL=/bin/zsh +CMD [ "zsh" ] From 9249b585101b44941203da2b6685d36d89d3c2ca Mon Sep 17 00:00:00 2001 From: Brian Ryu Date: Mon, 2 Feb 2026 19:11:22 +0000 Subject: [PATCH 2/5] Add missing dockerfile for 131 --- docker/Dockerfile.cu131 | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 docker/Dockerfile.cu131 diff --git a/docker/Dockerfile.cu131 b/docker/Dockerfile.cu131 new file mode 100644 index 0000000000..d33a296032 --- /dev/null +++ b/docker/Dockerfile.cu131 @@ -0,0 +1,31 @@ +FROM nvidia/cuda:13.1.1-devel-ubuntu24.04 + +ENV DEBIAN_FRONTEND=noninteractive + +# Update package lists and install system dependencies +RUN apt-get update && apt-get install -y \ + curl \ + git \ + wget + +# Install python +COPY docker/install/install_python.sh /install/install_python.sh +RUN bash /install/install_python.sh /opt/conda py312 + +# Set home directory +WORKDIR /workspace + +RUN echo "source activate py312" >> ~/.bashrc +ENV PATH="/opt/conda/bin:$PATH" +ENV PATH="/opt/conda/envs/py312/bin:$PATH" + +# Triton +ENV TRITON_PTXAS_PATH="/usr/local/cuda/bin/ptxas" + +# Install torch and other python packages +COPY requirements.txt /install/requirements.txt +COPY docker/install/install_python_packages.sh /install/install_python_packages.sh +RUN bash /install/install_python_packages.sh cu131 + +# Install mpi4py in the conda environment +RUN conda install -n py312 -y mpi4py mpich From 19ee57a0c8211d48a3bb09cd9b68418695b2f485 Mon Sep 17 00:00:00 2001 From: Brian Ryu Date: Mon, 2 Feb 2026 19:17:21 +0000 Subject: [PATCH 3/5] Update cuda-python version --- docker/install/install_python_packages.sh | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docker/install/install_python_packages.sh b/docker/install/install_python_packages.sh index 54f9095061..fcde28b6a3 100644 --- a/docker/install/install_python_packages.sh +++ b/docker/install/install_python_packages.sh @@ -29,7 +29,10 @@ pip3 install -r /install/requirements.txt pip3 install responses pytest scipy build cuda-python nvidia-nvshmem-cu12 # Install cudnn package based on CUDA version -if [[ "$CUDA_VERSION" == *"cu13"* ]]; then +if [[ "$CUDA_VERSION" == "cu131" ]]; then + pip3 install --upgrade cuda-python==13.1 + pip3 install "nvidia-cudnn-cu13>=9.14.0.64" +elif [[ "$CUDA_VERSION" == "cu130" ]]; then pip3 install --upgrade cuda-python==13.0 pip3 install "nvidia-cudnn-cu13>=9.14.0.64" else From 61e30af13d70747f2dd6ae1061de18f375f2f103 Mon Sep 17 00:00:00 2001 From: Brian Ryu Date: Mon, 2 Feb 2026 19:27:47 +0000 Subject: [PATCH 4/5] Wildcard versioning for cuda-python and update cudnn dependence for 13.1 --- docker/install/install_python_packages.sh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docker/install/install_python_packages.sh b/docker/install/install_python_packages.sh index fcde28b6a3..ca11445850 100644 --- a/docker/install/install_python_packages.sh +++ b/docker/install/install_python_packages.sh @@ -30,10 +30,10 @@ pip3 install responses pytest scipy build cuda-python nvidia-nvshmem-cu12 # Install cudnn package based on CUDA version if [[ "$CUDA_VERSION" == "cu131" ]]; then - pip3 install --upgrade cuda-python==13.1 - pip3 install "nvidia-cudnn-cu13>=9.14.0.64" + pip3 install --upgrade cuda-python==13.1.* + pip3 install "nvidia-cudnn-cu13>=9.18.1.3" elif [[ "$CUDA_VERSION" == "cu130" ]]; then - pip3 install --upgrade cuda-python==13.0 + pip3 install --upgrade cuda-python==13.0.* pip3 install "nvidia-cudnn-cu13>=9.14.0.64" else pip3 install --upgrade cuda-python==12.* From a37bdcaca3a310cdf02ef5d6c4000f31e9a93f7a Mon Sep 17 00:00:00 2001 From: Brian Ryu Date: Mon, 2 Feb 2026 19:32:48 +0000 Subject: [PATCH 5/5] Nightly torch 131 --- docker/Dockerfile.cu131 | 3 ++- docker/Dockerfile.cu131.dev | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.cu131 b/docker/Dockerfile.cu131 index d33a296032..10a643bfe6 100644 --- a/docker/Dockerfile.cu131 +++ b/docker/Dockerfile.cu131 @@ -25,7 +25,8 @@ ENV TRITON_PTXAS_PATH="/usr/local/cuda/bin/ptxas" # Install torch and other python packages COPY requirements.txt /install/requirements.txt COPY docker/install/install_python_packages.sh /install/install_python_packages.sh -RUN bash /install/install_python_packages.sh cu131 +# use nightly/cu131 temporarily and change to cu131 when torch releases stable version +RUN bash /install/install_python_packages.sh nightly/cu131 # Install mpi4py in the conda environment RUN conda install -n py312 -y mpi4py mpich diff --git a/docker/Dockerfile.cu131.dev b/docker/Dockerfile.cu131.dev index 8b7ea81e98..72e9873c28 100644 --- a/docker/Dockerfile.cu131.dev +++ b/docker/Dockerfile.cu131.dev @@ -47,7 +47,8 @@ ENV PATH="/home/$USERNAME/conda/envs/py312/bin:$PATH" # Install torch and other python packages COPY requirements.txt /install/requirements.txt COPY docker/install/install_python_packages.sh /install/install_python_packages.sh -RUN bash /install/install_python_packages.sh cu131 && pip3 install pre-commit +# use nightly/cu131 temporarily and change to cu131 when torch releases stable version +RUN bash /install/install_python_packages.sh nightly/cu131 && pip3 install pre-commit # Install mpi4py in the conda environment RUN conda install -n py312 -y mpi4py mpich