Skip to content

docker: Add CUDA 13.2 Docker containers #2843

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yongwww merged 1 commit into from
Mar 22, 2026
Merged

docker: Add CUDA 13.2 Docker containers #2843

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions .github/workflows/release-ci-docker.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ jobs:
needs: generate-tag
strategy:
matrix:
cuda: [cu126, cu128, cu129, cu130, cu131]
cuda: [cu126, cu128, cu129, cu130, cu132]
arch: [amd64, arm64]
steps:
- uses: actions/checkout@v4
Expand Down Expand Up @@ -74,7 +74,7 @@ jobs:
needs: [generate-tag, build]
strategy:
matrix:
cuda: [cu126, cu128, cu129, cu130, cu131]
cuda: [cu126, cu128, cu129, cu130, cu132]
steps:
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3
Expand Down Expand Up @@ -111,7 +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}
flashinfer/flashinfer-ci-cu132: ${DATE_SHA}
EOF

- name: Create Pull Request
Expand All @@ -128,7 +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 }}
- flashinfer/flashinfer-ci-cu132:${{ 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 }}
Expand Down
44 changes: 44 additions & 0 deletions docker/Dockerfile.cu132
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
FROM nvidia/cuda:13.2.0-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
Comment on lines +6 to +9
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot Mar 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

It's a Docker best practice to clean up the apt cache within the same RUN layer that apt-get update is called. This reduces the final image size. Please add && rm -rf /var/lib/apt/lists/* to this command.

RUN apt-get update && apt-get install -y \
    curl \
    git \
    wget \
    && rm -rf /var/lib/apt/lists/*


# 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"

# Set LD_LIBRARY_PATH to ensure pip-installed nvidia-cublas takes precedence over system libraries
ENV LD_LIBRARY_PATH="/opt/conda/envs/py312/lib/python3.12/site-packages/nvidia/cu13/lib/:$LD_LIBRARY_PATH"

# Triton
ENV TRITON_PTXAS_PATH="/usr/local/cuda/bin/ptxas"

# Install torch and other python packages
# use nightly/cu132 temporarily and change to cu132 when torch releases stable version
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 nightly/cu132

# Install tilelang and cuda-tile
RUN pip install tilelang cuda-tile

# Install mpi4py in the conda environment
RUN conda install -n py312 -y mpi4py mpich

# Configure pip for user-site installations (allows arbitrary users to install packages)
# This enables 'pip install --user' and 'pip install -e .' to work for any user
RUN mkdir -p /opt/pip-user && chmod 1777 /opt/pip-user
ENV PYTHONUSERBASE=/opt/pip-user
ENV PATH="/opt/pip-user/bin:$PATH"
77 changes: 77 additions & 0 deletions docker/Dockerfile.cu132.dev
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
FROM nvidia/cuda:13.2.0-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/*
Comment on lines +6 to +14
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot Mar 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

To optimize Docker layers and reduce build time, it's best to install all apt packages in a single RUN instruction. Consider adding sudo here, as it's installed in a separate layer later on. This will also allow removing a redundant apt-get update.

RUN apt-get update && apt-get install -y \
    curl \
    git \
    wget \
    clang-format \
    clangd-19 \
    vim \
    zsh \
    sudo \
    && 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/*
Comment on lines +22 to +29
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist bot Mar 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

Following the previous suggestion to install sudo earlier, this RUN instruction can be simplified to only handle user creation and sudoers configuration. This avoids an unnecessary apt-get update and apt-get install, making the Dockerfile more efficient.

RUN groupadd --gid $USER_GID $USERNAME \
    && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME \
    # [Optional] Add sudo support
    && echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \
    && chmod 0440 /etc/sudoers.d/$USERNAME


# 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
# use nightly/cu132 temporarily and change to cu132 when torch releases stable version
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 nightly/cu132 && pip3 install pre-commit
Comment on lines +43 to +51
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Mar 20, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Keep the dev image on the same CUDA library search order as the CI image.

docker/Dockerfile.cu132 adds the cu13 wheel libs to LD_LIBRARY_PATH, but this dev variant installs the same nightly/cu132 stack without that override. That can make local repro load a different CUDA user-space than CI.

Suggested diff 
 RUN echo "source activate py312" >> ~/.bashrc
 ENV PATH="/home/$USERNAME/conda/bin:$PATH"
 ENV PATH="/home/$USERNAME/conda/envs/py312/bin:$PATH"
+ENV LD_LIBRARY_PATH="/home/$USERNAME/conda/envs/py312/lib/python3.12/site-packages/nvidia/cu13/lib/:$LD_LIBRARY_PATH"
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@docker/Dockerfile.cu132.dev` around lines 43 - 51, The dev Dockerfile is
missing the CUDA wheel library path override, causing different CUDA user-space
ordering than CI; update the Dockerfile.cu132.dev to prepend the cu132 wheel
libs into LD_LIBRARY_PATH (e.g., set ENV LD_LIBRARY_PATH to include
/home/$USERNAME/conda/envs/py312/lib before $LD_LIBRARY_PATH) so the
nightly/cu132 Python wheel libs are searched first (make the change near the
existing ENV PATH settings or the RUN that installs nightly/cu132).


# Install tilelang and cuda-tile
RUN pip install tilelang cuda-tile

# 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" ]
Loading