Here's a Dockerfile to add sageattention to cu128-slim

[vugenti]

Dockerfile.cu128-slim-sageattention.txt

Great project, thanks so much for your effort! I had some trouble getting sageattention to install but eventually settled on this solution. It took me a few hours to get it working right so hopefully this will help someone else. This is probably easily adapted for other images. Don't forget to add --use-sage-attention to CLI_ARGS (or maybe that part is optional)!

# ============================================================
# Stage 1 — Torch base with CUDA 12.8.1, cuDNN, Python 3.12
# ============================================================
FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 AS pytorch-cuda-sdk

# Basic build deps (no Python reinstall)
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        python3-pip python3.12-dev build-essential ninja-build git && \
    pip install --break-system-packages numpy && \
    rm -rf /var/lib/apt/lists/*

# Cache torch install to speed up rebuilds
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install --break-system-packages --index-url https://download.pytorch.org/whl/cu128 \
    torch==2.9.0+cu128

# ============================================================
# Stage 2 — Build SageAttention wheel
# ============================================================
FROM pytorch-cuda-sdk AS builder

# CUDA build vars
ENV CUDA_HOME=/usr/local/cuda
ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 9.0"

RUN git clone https://github.com/thu-ml/SageAttention.git /src/sageattention
WORKDIR /src/sageattention

RUN pip install --break-system-packages packaging && \
    pip wheel . --no-build-isolation -w /wheelhouse

# ============================================================
# Stage 3 — Runtime: ComfyUI + SageAttention wheel
# ============================================================
FROM yanwk/comfyui-boot:cu128-slim

COPY --from=builder /wheelhouse /tmp/wheelhouse
RUN pip install --no-cache-dir /tmp/wheelhouse/sageattention*.whl && rm -rf /tmp/wheelhouse

[YanWenKun]

@vugenti Thank you for sharing with us!

I can't just put SageAttention in slim images since it could have compatibility issues for GPUs other than Ampere/Ada.
But I just built the wheels for SageAttention and SpargeAttention, and pre-installed them in the cu128-megapak-pt28 image:

https://github.com/YanWenKun/ComfyUI-Docker/tree/main/cu128-megapak-pt28