From c04e4b62504b3216d627405b47a5c59d23960a22 Mon Sep 17 00:00:00 2001 From: Debian Date: Sat, 3 Jan 2026 22:01:58 +0000 Subject: [PATCH] Update for RTX 5090 (Blackwell sm_120) support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Switch to PyTorch nightly with CUDA 12.8 (required for sm_120) - Target TORCH_CUDA_ARCH_LIST="12.0" for Blackwell - Remove nunchaku (incompatible with PyTorch nightly) - Use latest SageAttention (has sm_120 kernel support) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- Dockerfile | 30 ++++++++++++------------------ 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/Dockerfile b/Dockerfile index 43bfe71..bae48c3 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -# ComfyUI RunPod Serverless - CUDA 12.8.1, Python 3.12, PyTorch 2.8.0 +# ComfyUI RunPod Serverless - CUDA 12.8.1, Python 3.12, PyTorch nightly (RTX 5090/Blackwell) FROM nvidia/cuda:12.8.1-devel-ubuntu22.04 ENV DEBIAN_FRONTEND=noninteractive @@ -44,34 +44,30 @@ RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.12 1 & python3.12 -m ensurepip --upgrade && \ python3.12 -m pip install --upgrade pip setuptools wheel -# Install PyTorch 2.8.0+cu128 and triton 3.4.0 in single layer -RUN pip install \ - torch==2.8.0+cu128 \ - torchvision==0.23.0+cu128 \ - torchaudio==2.8.0+cu128 \ - triton==3.4.0 \ - --index-url https://download.pytorch.org/whl/cu128 && \ +# Install PyTorch nightly with CUDA 12.8 (required for RTX 5090/Blackwell sm_120) +RUN pip install --pre \ + torch \ + torchvision \ + torchaudio \ + --index-url https://download.pytorch.org/whl/nightly/cu128 && \ rm -rf /root/.cache/pip /tmp/* -# Install nunchaku from GitHub wheel -RUN pip install https://github.com/nunchaku-tech/nunchaku/releases/download/v1.0.2/nunchaku-1.0.2+torch2.8-cp312-cp312-linux_x86_64.whl +# Note: nunchaku disabled - requires stable PyTorch, not compatible with nightly +# RUN pip install https://github.com/nunchaku-tech/nunchaku/releases/download/v1.0.2/nunchaku-1.0.2+torch2.8-cp312-cp312-linux_x86_64.whl # Install key dependencies before SageAttention COPY requirements.txt /tmp/requirements.txt RUN pip install -r /tmp/requirements.txt && rm -rf /root/.cache/pip -# Compile SageAttention 2.2 from source with no build isolation +# Compile SageAttention from source for RTX 5090 (Blackwell sm_120) WORKDIR /tmp ENV EXT_PARALLEL=2 ENV NVCC_APPEND_FLAGS="--threads 2" ENV MAX_JOBS=4 -# Target RunPod GPU architectures: -# 8.0 = A100, 9.0 = H100/H200 -# Note: 8.6/8.9 excluded due to SageAttention sm90 kernel cross-compile issues -ENV TORCH_CUDA_ARCH_LIST="8.0;9.0" +# Target RTX 5090 (Blackwell) - sm_120 = compute capability 12.0 +ENV TORCH_CUDA_ARCH_LIST="12.0" RUN git clone https://github.com/thu-ml/SageAttention.git && \ cd SageAttention && \ - git checkout 2aecfa8 && \ pip install --no-build-isolation . && \ cd / && rm -rf /tmp/SageAttention /root/.cache/pip @@ -93,7 +89,6 @@ RUN git clone https://github.com/ltdrdata/ComfyUI-Manager.git && \ git clone https://github.com/willmiao/ComfyUI-Lora-Manager.git && \ git clone https://github.com/city96/ComfyUI-GGUF.git && \ git clone https://github.com/Fannovel16/ComfyUI-Frame-Interpolation.git && \ - git clone https://github.com/nunchaku-tech/ComfyUI-nunchaku.git && \ git clone https://github.com/evanspearman/ComfyMath.git && \ git clone https://github.com/ssitu/ComfyUI_UltimateSDUpscale.git && \ git clone https://github.com/pythongosssss/ComfyUI-Custom-Scripts.git && \ @@ -106,7 +101,6 @@ RUN (cd ComfyUI-KJNodes && pip install -r requirements.txt || true) && \ (cd ComfyUI-VideoHelperSuite && pip install -r requirements.txt || true) && \ (cd ComfyUI-GGUF && pip install -r requirements.txt || true) && \ (cd ComfyUI-Frame-Interpolation && pip install -r requirements.txt || true) && \ - (cd ComfyUI-nunchaku && pip install -r requirements.txt || true) && \ (cd ComfyUI-Easy-Use && pip install -r requirements.txt || true) && \ rm -rf /root/.cache/pip /tmp/*