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 b/docker/Dockerfile.cu131 new file mode 100644 index 0000000000..10a643bfe6 --- /dev/null +++ b/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 diff --git a/docker/Dockerfile.cu131.dev b/docker/Dockerfile.cu131.dev new file mode 100644 index 0000000000..72e9873c28 --- /dev/null +++ b/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" ] diff --git a/docker/install/install_python_packages.sh b/docker/install/install_python_packages.sh index 54f9095061..ca11445850 100644 --- a/docker/install/install_python_packages.sh +++ b/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.*