Skip to content
Open
Changes from all commits
Commits
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
75 changes: 75 additions & 0 deletions Dockerfile.blackwell
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04

# Blackwell (RTX 50xx) requires sm_120 architecture
ARG CUDA_ARCHITECTURES="12.0"

ENV DEBIAN_FRONTEND=noninteractive

# Install system dependencies
RUN apt update && \
apt install -y \
python3 python3-pip git wget curl cmake ninja-build \
libgl1 libglib2.0-0 ffmpeg && \
apt clean

WORKDIR /workspace

COPY requirements.txt .

# Upgrade pip first
RUN pip install --upgrade pip setuptools wheel

# Install PyTorch 2.7.1 with CUDA 12.8 for Blackwell
RUN pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu128

# Install requirements
RUN pip install -r requirements.txt

# Install SageAttention from git (patch GPU detection for Blackwell)
ENV TORCH_CUDA_ARCH_LIST="${CUDA_ARCHITECTURES}"
ENV FORCE_CUDA="1"
ENV MAX_JOBS="1"

COPY <<EOF /tmp/patch_setup.py
import os
with open('setup.py', 'r') as f:
content = f.read()

# Get architectures from environment variable
arch_list = os.environ.get('TORCH_CUDA_ARCH_LIST')
arch_set = '{' + ', '.join([f'"{arch}"' for arch in arch_list.split(';')]) + '}'

# Replace the GPU detection section
old_section = '''compute_capabilities = set()
device_count = torch.cuda.device_count()
for i in range(device_count):
major, minor = torch.cuda.get_device_capability(i)
if major < 8:
warnings.warn(f"skipping GPU {i} with compute capability {major}.{minor}")
continue
compute_capabilities.add(f"{major}.{minor}")'''

new_section = 'compute_capabilities = ' + arch_set + '''
print(f"Manually set compute capabilities: {compute_capabilities}")'''

content = content.replace(old_section, new_section)

with open('setup.py', 'w') as f:
f.write(content)
EOF

RUN git clone https://github.com/thu-ml/SageAttention.git /tmp/sageattention && \
cd /tmp/sageattention && \
python3 /tmp/patch_setup.py && \
pip install --no-build-isolation .

RUN useradd -u 1000 -ms /bin/bash user

RUN chown -R user:user /workspace

RUN mkdir /home/user/.cache && \
chown -R user:user /home/user/.cache

COPY entrypoint.sh /workspace/entrypoint.sh

ENTRYPOINT ["/workspace/entrypoint.sh"]