diff --git a/apps/dreamverse/docker/Dockerfile b/apps/dreamverse/docker/Dockerfile index c47d20f82..680dd1c2a 100644 --- a/apps/dreamverse/docker/Dockerfile +++ b/apps/dreamverse/docker/Dockerfile @@ -47,7 +47,11 @@ COPY . /opt/FastVideo RUN --mount=type=cache,target=/opt/uv/cache \ source /opt/venv/bin/activate \ - && uv pip install "/opt/FastVideo[dreamverse]" + && if [[ "${UV_TORCH_BACKEND}" == "cu130" ]]; then \ + uv pip install "/opt/FastVideo[dreamverse]" "nvidia-cutlass-dsl[cu13]"; \ + else \ + uv pip install "/opt/FastVideo[dreamverse]"; \ + fi # Standard docker build does not expose GPUs, while fastvideo-kernel/build.sh # detects the CUDA architecture with torch at build time. The FastVideo package diff --git a/docker/Dockerfile b/docker/Dockerfile index 26fe8e79e..7ed63e26f 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -66,7 +66,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p # 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 CuTe DSL 4.6 that -# flashinfer/quack pull in. After the wheel install we overlay this 4.6-compatible +# flashinfer/quack pull in. After the wheel install we overlay this 4.6.0.dev0-compatible # upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback. ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083 @@ -169,18 +169,24 @@ RUN --mount=type=cache,target=/opt/uv/cache \ # 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. x86 only: the FA4 stack (quack-kernels etc.) is unvalidated on -# arm64 / GB10 (sm_121), so there we skip the overlay; FA4 is opt-in -# (FASTVIDEO_FA4=1) and errors if set without the overlay, so leave it unset on -# arm64 and the image runs FA3/FA2 as usual. +# an FA2-less image. CUDA 13 selects the matching runtime extra on both amd64 and +# arm64 (validated on GB200). CUDA 12.6 amd64 keeps the bare/default runtime; +# CUDA 12.6 arm64 remains unvalidated and skips the overlay. FA4 stays opt-in, so +# installing it does not auto-enable the backend on GB10 (sm_121). RUN --mount=type=cache,target=/opt/uv/cache \ source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ - if [ "${TARGETARCH}" = "arm64" ]; then \ - echo "Skipping FA4 cute overlay on arm64 (FA4 stack unvalidated there; do not set FASTVIDEO_FA4)"; \ + if [ "${UV_TORCH_BACKEND}" = "cu130" ]; then \ + FA4_PACKAGE="flash-attn-4[cu13]"; \ + elif [ "${TARGETARCH}" = "arm64" ]; then \ + echo "Skipping FA4 cute overlay on CUDA 12.6 arm64 (do not set FASTVIDEO_FA4)"; \ + FA4_PACKAGE=""; \ else \ + FA4_PACKAGE="flash-attn-4"; \ + fi && \ + if [ -n "${FA4_PACKAGE}" ]; then \ 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')]" && \ - uv pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \ + uv pip install "${FA4_PACKAGE} @ 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')"; \ fi diff --git a/docs/inference/optimizations.md b/docs/inference/optimizations.md index a11fccf1c..c14be2138 100644 --- a/docs/inference/optimizations.md +++ b/docs/inference/optimizations.md @@ -78,8 +78,8 @@ python setup.py install FastVideo never auto-selects FlashAttention-4 (`flash_attn.cute`) just because it is installed: its CuTeDSL kernels JIT-compile per shape family and can fail at -runtime on some GPU/shape combinations. To use FA4, install the pinned -`flash-attn-4` build (see the `flash-attn-4` source in `pyproject.toml`) and set: +runtime on some GPU/shape combinations. To use FA4 on CUDA 13, install the pinned +`flash-attn-4[cu13]` build (see the `flash-attn-4` source in `pyproject.toml`) and set: ```bash export FASTVIDEO_FA4=1 diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 0bf0c1f77..8df137bb1 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -81,7 +81,7 @@ have an incompatible `_flash_attn_fwd` signature. ```bash pip install torchvision -pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@82d6441eec5d4dfec120153db2c0145ae855a083#subdirectory=flash_attn/cute" +pip install "flash-attn-4[cu13] @ git+https://github.com/Dao-AILab/flash-attention.git@82d6441eec5d4dfec120153db2c0145ae855a083#subdirectory=flash_attn/cute" ``` The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against diff --git a/fastvideo/tests/contract/test_fa4_dependency_pins.py b/fastvideo/tests/contract/test_fa4_dependency_pins.py index a18c76818..8319232c1 100644 --- a/fastvideo/tests/contract/test_fa4_dependency_pins.py +++ b/fastvideo/tests/contract/test_fa4_dependency_pins.py @@ -30,3 +30,21 @@ def test_fa4_dependency_pins_match() -> None: } assert len(set(pins.values())) == 1, f"FA4 dependency pins disagree: {pins}" + + +def test_fa4_blackwell_installs_select_cuda13_runtime() -> None: + root_pyproject = (REPO_ROOT / "pyproject.toml").read_text(encoding="utf-8") + dreamverse_pyproject = (REPO_ROOT / "apps/dreamverse/pyproject.toml").read_text(encoding="utf-8") + dockerfile = (REPO_ROOT / "docker/Dockerfile").read_text(encoding="utf-8") + dreamverse_dockerfile = (REPO_ROOT / "apps/dreamverse/docker/Dockerfile").read_text(encoding="utf-8") + kernel_readme = (REPO_ROOT / "fastvideo-kernel/README.md").read_text(encoding="utf-8") + + assert '"flash-attn-4",' in root_pyproject + assert '"flash-attn-4 @ git+' in dreamverse_pyproject + assert 'if [ "${UV_TORCH_BACKEND}" = "cu130" ]' in dockerfile + assert 'FA4_PACKAGE="flash-attn-4[cu13]"' in dockerfile + assert 'FA4_PACKAGE="flash-attn-4"' in dockerfile + assert 'elif [ "${TARGETARCH}" = "arm64" ]' in dockerfile + assert '"${UV_TORCH_BACKEND}" == "cu130"' in dreamverse_dockerfile + assert '"nvidia-cutlass-dsl[cu13]"' in dreamverse_dockerfile + assert "flash-attn-4[cu13]" in kernel_readme diff --git a/pyproject.toml b/pyproject.toml index 99f96c0c7..a3089a902 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -114,7 +114,7 @@ fastvideo-kernel = [ # uv reads UV_TORCH_BACKEND / --torch-backend (not a pyproject setting). The # supported backends (cu126, cu130) both provide torch 2.12.0. imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" } -# FA4 cute, pinned to a CuTe DSL 4.6-compatible revision. torch.compile support +# FA4 cute, pinned to the CuTe DSL 4.6.0.dev0-compatible revision. torch.compile support # comes from FastVideo's own custom_op wrappers. flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "82d6441eec5d4dfec120153db2c0145ae855a083", subdirectory = "flash_attn/cute" }