ci: Add CUDA 13.1 CI container support

.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"
+}

.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 }}

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

docker/Dockerfile.cu131

@@ -0,0 +1,32 @@
+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
+# 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

docker/Dockerfile.cu131.dev

@@ -0,0 +1,74 @@
+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
+# 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
+
+# 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" ]

docker/install/install_python_packages.sh

@@ -29,8 +29,11 @@ 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
-  pip3 install --upgrade cuda-python==13.0
+if [[ "$CUDA_VERSION" == "cu131" ]]; then
+  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 "nvidia-cudnn-cu13>=9.14.0.64"
 else
   pip3 install --upgrade cuda-python==12.*