SolitaryThinker
fba50356b1
FA4 cute supports SM 9/10/11 only, and the pinned revision forces nvidia-cutlass-dsl==4.6.0.dev0 over the 4.8.x that flashinfer 0.7.2 needs, so import flashinfer fails in the image. GB10 FP4 attention runs on fastvideo-kernel. Add INSTALL_FA4_CUTE (default 1), set it to 0 for the CUDA 13 arm64 leg, still remove the wheel's stale cute files, and check that flashinfer imports when the overlay is skipped.
262 lines
13 KiB
Docker
262 lines
13 KiB
Docker
# syntax=docker/dockerfile:1.7
|
|
# =============================================================================
|
|
# FastVideo CUDA development image (unified, parameterized)
|
|
# =============================================================================
|
|
# Replaces the former per-Python Dockerfiles
|
|
# (Dockerfile.python3.10 / 3.11 / 3.12 / 3.12.cuda12.9.1): pick a variant with
|
|
# --build-arg instead of maintaining a separate near-identical file.
|
|
#
|
|
# # Default image (Python 3.12, CUDA 13.0.0, Ubuntu 22.04, Hopper sm_90a):
|
|
# docker build -f docker/Dockerfile -t fastvideo-dev:py3.12 .
|
|
#
|
|
# # Python 3.10 / 3.11:
|
|
# docker build -f docker/Dockerfile --build-arg PYTHON_VERSION=3.11 \
|
|
# -t fastvideo-dev:py3.11 .
|
|
#
|
|
# # CUDA 13.0.0: match the torch backend + flash-attn wheel to the base CUDA.
|
|
# docker build -f docker/Dockerfile \
|
|
# --build-arg CUDA_VERSION=13.0.0 \
|
|
# --build-arg UV_TORCH_BACKEND=cu130 \
|
|
# --build-arg FLASH_ATTN_WHEEL_TAG=cu130torch2.12 \
|
|
# -t fastvideo-dev:py3.12-cuda13.0.0 .
|
|
#
|
|
# # DGX Spark / GB10 (aarch64, Blackwell sm_121a): the matching arm64
|
|
# # prebuilt flash-attn wheel is selected automatically from TARGETARCH.
|
|
# docker build --platform linux/arm64 -f docker/Dockerfile \
|
|
# --build-arg TORCH_CUDA_ARCH_LIST=12.1a --build-arg INSTALL_FA4_CUTE=0 \
|
|
# -t fastvideo-dev:spark .
|
|
#
|
|
# The ARG defaults below are the source of truth for the default image.
|
|
# ROCm builds use docker/rocm7_1.Dockerfile (different base + flow).
|
|
# =============================================================================
|
|
|
|
# --- Base image selection (overridable for private mirrors / new CUDA) -------
|
|
# Build on the -devel- image: nvcc + the full toolkit are needed to compile
|
|
# fastvideo-kernel. CUDA_VERSION/UBUNTU_VERSION feed the default tag, or override
|
|
# BUILD_BASE_IMAGE wholesale for a mirror.
|
|
ARG CUDA_VERSION=13.0.0
|
|
ARG UBUNTU_VERSION=22.04
|
|
ARG BUILD_BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu${UBUNTU_VERSION}
|
|
|
|
FROM ${BUILD_BASE_IMAGE} AS base
|
|
|
|
# Redeclare the ARGs we need inside the stage (FROM consumes the global ones).
|
|
ARG CUDA_VERSION
|
|
ARG PYTHON_VERSION=3.12
|
|
|
|
# CUDA arch fastvideo-kernel is compiled for. Build machines have no GPU, so
|
|
# build.sh uses this instead of probing a live device. Default Hopper (sm_90a),
|
|
# matching the released kernel wheel; override e.g. TORCH_CUDA_ARCH_LIST=12.1a
|
|
# for GB10/Blackwell (DGX Spark).
|
|
ARG TORCH_CUDA_ARCH_LIST=9.0a
|
|
|
|
# Cap CMake/Ninja concurrency to fit standard 16 GB build runners. This default
|
|
# matches the CUDA 13 kernel-release build; the CUDA 12.6 image workflow uses 4.
|
|
ARG CMAKE_BUILD_PARALLEL_LEVEL=1
|
|
|
|
# flash-attn: when FLASH_ATTN_WHEEL_TAG is set, install the matching prebuilt wheel
|
|
# for TARGETARCH. The prebuilt-wheel project publishes x86_64 and arm64 wheels in
|
|
# different releases, so each architecture has its own release URL. Set the wheel
|
|
# tag to empty only when an explicit source build is desired. The wheel's
|
|
# cuXXXtorchY.Y tag must match UV_TORCH_BACKEND (below): cu130torch2.12 for the
|
|
# default 13.0 base, cu126torch2.12 for 12.6 (both backends are on torch 2.12).
|
|
ARG FLASH_ATTN_VERSION=2.8.3
|
|
ARG FLASH_ATTN_WHEEL_TAG=cu130torch2.12
|
|
ARG FLASH_ATTN_WHEEL_RELEASE=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.17
|
|
ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.22
|
|
|
|
# The prebuilt flash-attn wheel ships an FA4 `flash_attn.cute` built against the
|
|
# cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
|
|
# flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
|
|
# upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
|
|
ARG FA4_CUTE_REF=14c377950125c70b7a9dabf9c561fca53715ac7d
|
|
# FA4 cute supports SM 9/10/11 only; GB10 (sm_121) runs FP4 attention through
|
|
# fastvideo-kernel instead. INSTALL_FA4_CUTE=0 skips the overlay so the image
|
|
# keeps the cutlass-dsl that flashinfer resolves: this revision pins
|
|
# nvidia-cutlass-dsl==4.6.0.dev0, which flashinfer 0.7.2 cannot import against.
|
|
ARG INSTALL_FA4_CUTE=1
|
|
|
|
# Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
|
|
# select the prebuilt flash-attn wheel. Empty under a plain `docker build` without
|
|
# BuildKit, which we treat as x86_64 for compatibility with the historical build.
|
|
ARG TARGETARCH
|
|
|
|
ENV DEBIAN_FRONTEND=noninteractive
|
|
|
|
SHELL ["/bin/bash", "-c"]
|
|
|
|
WORKDIR /FastVideo
|
|
|
|
RUN apt-get update && apt-get install -y --no-install-recommends \
|
|
wget \
|
|
git \
|
|
ca-certificates \
|
|
openssh-server \
|
|
zsh \
|
|
vim \
|
|
curl \
|
|
ffmpeg \
|
|
libgl1 \
|
|
libglib2.0-0 \
|
|
libx11-dev \
|
|
gcc-11 \
|
|
g++-11 \
|
|
clang-11 \
|
|
cmake \
|
|
pkg-config \
|
|
build-essential \
|
|
libssl-dev \
|
|
&& rm -rf /var/lib/apt/lists/*
|
|
|
|
# Rust toolchain: some dependencies only ship sdists on aarch64 and need cargo
|
|
# to build. The dormant legacy Modal image layers the identical apt set +
|
|
# rustup on top of this image (fastvideo/tests/modal/pr_test.py); baking both
|
|
# here keeps its manual rollback path reproducible without changing the Slurm
|
|
# runner's package surface.
|
|
RUN set -o pipefail && \
|
|
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable --profile minimal && \
|
|
/root/.cargo/bin/cargo --version && /root/.cargo/bin/rustc --version
|
|
ENV PATH=/root/.cargo/bin:${PATH}
|
|
|
|
# Set up C++20 compilers for ThunderKittens
|
|
RUN update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 100 --slave /usr/bin/g++ g++ /usr/bin/g++-11
|
|
|
|
# Set CUDA environment variables. Use the version-agnostic /usr/local/cuda
|
|
# symlink the nvidia images provide so this works for any CUDA_VERSION.
|
|
ENV CUDA_HOME=/usr/local/cuda
|
|
ENV PATH=${CUDA_HOME}/bin:${PATH}
|
|
ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:$LD_LIBRARY_PATH
|
|
|
|
# The uv package cache lives in a BuildKit cache mount (see RUNs below), not in an
|
|
# image layer, so we can drop the old --no-cache-dir to speed up rebuilds without
|
|
# bloating the image. Copy link mode avoids hardlink errors across the mount.
|
|
ENV UV_CACHE_DIR=/opt/uv/cache
|
|
ENV UV_LINK_MODE=copy
|
|
|
|
# pyproject no longer pins a PyTorch index; select the build at resolve time.
|
|
# Build hosts have no GPU driver (auto-detect would fall back to CPU), so pin
|
|
# explicitly, matching the base CUDA: cu130 (default, 13.0) / cu126 (12.6), both
|
|
# on torch 2.12. Keep FLASH_ATTN_WHEEL_TAG in sync.
|
|
ARG UV_TORCH_BACKEND=cu130
|
|
ENV UV_TORCH_BACKEND=${UV_TORCH_BACKEND}
|
|
|
|
# Install uv and source its environment
|
|
RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
|
|
echo 'source $HOME/.local/bin/env' >> /root/.bashrc
|
|
|
|
# Copy only dependency metadata first to leverage Docker cache. The kernel is
|
|
# deliberately excluded from this layer: arm64 resolves it to the in-tree source,
|
|
# which needs the complete checkout and CUDA architecture settings below.
|
|
COPY pyproject.toml ./
|
|
COPY docker/uv-excludes ./docker/uv-excludes
|
|
|
|
# Create a dummy README to satisfy the installation
|
|
RUN echo "# Placeholder" > README.md
|
|
|
|
# Create the venv and install the (slow-changing) dependency tree + flash-attn
|
|
# from pyproject.toml alone, before the source is copied in, so that editing
|
|
# source below does not bust this expensive layer.
|
|
RUN --mount=type=cache,target=/opt/uv/cache \
|
|
source $HOME/.local/bin/env && \
|
|
uv venv --python ${PYTHON_VERSION} --seed /opt/venv && \
|
|
source /opt/venv/bin/activate && \
|
|
uv pip install --upgrade pip && \
|
|
uv pip install --excludes docker/uv-excludes ".[dev]" && \
|
|
python -c "import cv2; print('OpenCV', cv2.__version__)" && \
|
|
PYTAG=cp$(echo "${PYTHON_VERSION}" | tr -d .) && \
|
|
case "${TARGETARCH:-amd64}" in \
|
|
amd64) \
|
|
FLASH_ATTN_WHEEL_PLATFORM=linux_x86_64; \
|
|
FLASH_ATTN_TARGET_RELEASE=${FLASH_ATTN_WHEEL_RELEASE} \
|
|
;; \
|
|
arm64) \
|
|
FLASH_ATTN_WHEEL_PLATFORM=linux_aarch64; \
|
|
FLASH_ATTN_TARGET_RELEASE=${FLASH_ATTN_WHEEL_RELEASE_ARM64} \
|
|
;; \
|
|
*) \
|
|
echo "Unsupported TARGETARCH for flash-attn: ${TARGETARCH}" >&2; \
|
|
exit 1 \
|
|
;; \
|
|
esac && \
|
|
if [ -n "${FLASH_ATTN_WHEEL_TAG}" ]; then \
|
|
echo "Installing prebuilt flash-attn wheel (${FLASH_ATTN_WHEEL_TAG}, ${PYTAG}, ${FLASH_ATTN_WHEEL_PLATFORM})" && \
|
|
uv pip install "${FLASH_ATTN_TARGET_RELEASE}/flash_attn-${FLASH_ATTN_VERSION}+${FLASH_ATTN_WHEEL_TAG}-${PYTAG}-${PYTAG}-${FLASH_ATTN_WHEEL_PLATFORM}.whl"; \
|
|
else \
|
|
echo "Building flash-attn ${FLASH_ATTN_VERSION} from source" && \
|
|
uv pip install flash-attn==${FLASH_ATTN_VERSION} --no-build-isolation; \
|
|
fi
|
|
|
|
# Overlay the CuTe-DSL-4.6-compatible upstream FA4 cute (FA4_CUTE_REF) over the
|
|
# wheel/source one so the image runs FA4, not the FA2 fallback. This pulls the FA4
|
|
# runtime stack (cutlass-dsl, quack-kernels, apache-tvm-ffi, torch-c-dlpack-ext) --
|
|
# the same deps the [dreamverse] extra already installs in CI; the installed torch
|
|
# is left untouched. The cute package only ships flash_attn/cute (never
|
|
# flash_attn/__init__.py), so FA2/varlen/bert_padding stay from the install above;
|
|
# rmtree clears the wheel's stale cute files first to avoid an install conflict.
|
|
# Then verify both survive so a broken overlay fails the build instead of shipping
|
|
# an FA2-less image. The pinned stack is validated on ARM64 GB200 (sm_100) as well
|
|
# as x86; FA4 remains opt-in through FASTVIDEO_FA4=1 so lanes with FA2 baselines
|
|
# keep their existing numerics.
|
|
# With INSTALL_FA4_CUTE=0 the wheel's stale cute is still removed, so
|
|
# `import flash_attn.cute` fails cleanly and FastVideo's guarded import falls back.
|
|
RUN --mount=type=cache,target=/opt/uv/cache \
|
|
source $HOME/.local/bin/env && \
|
|
source /opt/venv/bin/activate && \
|
|
python -c "import glob, shutil; [shutil.rmtree(d, ignore_errors=True) for d in glob.glob('/opt/venv/lib/python*/site-packages/flash_attn/cute')]" && \
|
|
if [ "${INSTALL_FA4_CUTE}" = "1" ]; then \
|
|
uv pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \
|
|
python -c "import flash_attn; assert hasattr(flash_attn, 'flash_attn_func'), 'FA2 was clobbered by the cute overlay'; import flash_attn.cute; print('FA2 + FA4 cute OK')"; \
|
|
else \
|
|
python -c "import flash_attn, flashinfer; assert hasattr(flash_attn, 'flash_attn_func'); print('FA2 + flashinfer', flashinfer.__version__, 'OK; FA4 cute skipped')"; \
|
|
fi
|
|
|
|
COPY . .
|
|
|
|
# Build immutable FastVideo kernel wheels for the published image. The requested
|
|
# architecture remains installed for normal image users; amd64 images also carry
|
|
# an SM89 artifact for L40S users and the dormant legacy rollback path.
|
|
ARG FASTVIDEO_KERNEL_PREBUILT_DIR=/opt/fastvideo-kernel-prebuilt
|
|
RUN --mount=type=cache,target=/opt/uv/cache \
|
|
source $HOME/.local/bin/env && \
|
|
source /opt/venv/bin/activate && \
|
|
default_arch="${TORCH_CUDA_ARCH_LIST}" && \
|
|
default_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/${default_arch}" && \
|
|
export TORCH_CUDA_ARCH_LIST="${default_arch}" && \
|
|
cd fastvideo-kernel && \
|
|
CMAKE_ARGS= CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
|
|
./build.sh --wheel-dir "${default_wheel_dir}" && \
|
|
cd /FastVideo && \
|
|
CMAKE_ARGS= python fastvideo/tests/modal/kernel_build_cache.py write-build-info \
|
|
--wheel-dir "${default_wheel_dir}" \
|
|
--output "${default_wheel_dir}/metadata.json" && \
|
|
if [[ "${TARGETARCH:-amd64}" == "amd64" && "${default_arch}" != "8.9" ]]; then \
|
|
export TORCH_CUDA_ARCH_LIST=8.9 && \
|
|
l40s_wheel_dir="${FASTVIDEO_KERNEL_PREBUILT_DIR}/8.9" && \
|
|
cd /FastVideo/fastvideo-kernel && \
|
|
CMAKE_ARGS= CMAKE_BUILD_PARALLEL_LEVEL=${CMAKE_BUILD_PARALLEL_LEVEL} \
|
|
./build.sh --wheel-dir "${l40s_wheel_dir}" && \
|
|
cd /FastVideo && \
|
|
CMAKE_ARGS= python fastvideo/tests/modal/kernel_build_cache.py write-build-info \
|
|
--wheel-dir "${l40s_wheel_dir}" \
|
|
--output "${l40s_wheel_dir}/metadata.json"; \
|
|
fi && \
|
|
export TORCH_CUDA_ARCH_LIST="${default_arch}" && \
|
|
default_wheel="$(find "${default_wheel_dir}" -maxdepth 1 -type f \
|
|
\( -name 'fastvideo_kernel-*.whl' -o -name 'fastvideo-kernel-*.whl' \) \
|
|
| sort | tail -n 1)" && \
|
|
uv pip install "${default_wheel}" \
|
|
--reinstall-package fastvideo-kernel --no-deps
|
|
|
|
# Install FastVideo itself (editable) now that the source is present, and set up
|
|
# shell configuration. Dependencies and the local kernel are already installed,
|
|
# so --no-deps prevents uv from resolving/building the ARM source override again.
|
|
RUN --mount=type=cache,target=/opt/uv/cache \
|
|
source $HOME/.local/bin/env && \
|
|
source /opt/venv/bin/activate && \
|
|
uv pip install -e ".[dev]" --no-deps && \
|
|
(git config --unset-all http.https://github.com/.extraheader || true) && \
|
|
echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
|
|
echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
|
|
|
|
EXPOSE 22
|