[bugfix]: install the FA4 CUDA 13 runtime

Select flash-attn-4's cu13 extra on Blackwell install surfaces so revision 82d6441 receives its exact CuTe DSL 4.6.0.dev0 runtime. Enable the opt-in overlay for arm64 CUDA 13 images, including GB200, while CUDA 12.6 continues to use FA3/FA2.

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

apps/dreamverse/pyproject.toml

@@ -14,7 +14,7 @@ dependencies = [
 [project.optional-dependencies]
 server = [
     "cerebras-cloud-sdk",
-    "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute",
+    "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@82d6441eec5d4dfec120153db2c0145ae855a083#subdirectory=flash_attn/cute",
     "flashinfer-python",
     "openai>=1.40",
 ]

docker/Dockerfile

@@ -65,8 +65,8 @@ ARG FLASH_ATTN_WHEEL_RELEASE=https://github.com/mjun0812/flash-attention-prebuil
 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
+# 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.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
 

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
@@ -110,14 +110,18 @@ See the [Attn-QAT paper](https://arxiv.org/abs/2603.00040) and [flash-attention-
 Install the FP4 flash attention kernel (without upgrading your existing torch):
 
 ```bash
-# branch fix/cutlass-dsl-4.5 carries the cutlass-dsl 4.5 fix (cute.core.ThrMma
-# -> cute.ThrMma); switch back to @fp4 once hao-ai-lab/flash-attention-fp4#2 merges.
-pip install --no-deps "git+ssh://git@github.com/hao-ai-lab/flash-attention-fp4.git@fix/cutlass-dsl-4.5#subdirectory=flash_attn/cute"
-pip install "nvidia-cutlass-dsl>=4.5.2" apache-tvm-ffi flashinfer-python
+# Commit 940bf7e5 carries the CuTe DSL 4.5 fix (cute.core.ThrMma -> cute.ThrMma).
+pip install --no-deps "git+https://github.com/hao-ai-lab/flash-attention-fp4.git@940bf7e511375ec160bc2d7188bef35915ded1e3#subdirectory=flash_attn/cute"
+pip install "nvidia-cutlass-dsl[cu13]==4.5.2" "quack-kernels==0.5.0" apache-tvm-ffi flashinfer-python torch-c-dlpack-ext
 ```
 
 The `--no-deps` flag prevents upgrading torch/torchvision. Use the supported
-PyTorch 2.12.0 and CUDA 13 environment for this kernel.
+PyTorch 2.12.0 and CUDA 13 environment for this kernel. Keep this private FP4
+overlay in a separate inference environment from FastVideo's upstream dense FA4
+dependency: the two distributions provide the same `flash_attn.cute` package but
+require different CuTe DSL versions. Keep the CUTLASS and Quack pins together;
+Quack 0.5.1 and newer use CuTe DSL 4.6, whose API is incompatible with this FP4
+kernel revision.
 
 #### Usage
 

fastvideo-kernel/README.md

@@ -75,13 +75,13 @@ fully usable without it).
 The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`,
 `flash_attn.cute.interface._flash_attn_fwd`) are provided upstream by
 [Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention). Pin to
-commit `940cd9680f3315f2f06b43ab5bea2c2cf2d96806`, the revision FastVideo pins as
+commit `82d6441eec5d4dfec120153db2c0145ae855a083`, the revision FastVideo pins as
 the `flash-attn-4` source in the repo-root `pyproject.toml`; other revisions may
 have an incompatible `_flash_attn_fwd` signature.
 
 ```bash
-pip install "nvidia-cutlass-dsl>=4.5.0" torchvision
-pip install "git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute"
+pip install torchvision
+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

fastvideo/tests/contract/test_fa4_dependency_pins.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Keep every FA4 installation surface on the same tested revision."""
+
+import re
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+
+
+def _match(path: str, pattern: str) -> str:
+    source = (REPO_ROOT / path).read_text(encoding="utf-8")
+    match = re.search(pattern, source)
+    assert match is not None, f"FA4 pin not found in {path}"
+    return match.group(1)
+
+
+def test_fa4_dependency_pins_match() -> None:
+    pins = {
+        "pyproject.toml": _match(
+            "pyproject.toml",
+            r'flash-attn-4 = .*rev = "([0-9a-f]{40})"'),
+        "docker/Dockerfile": _match(
+            "docker/Dockerfile", r"ARG FA4_CUTE_REF=([0-9a-f]{40})"),
+        "apps/dreamverse/pyproject.toml": _match(
+            "apps/dreamverse/pyproject.toml",
+            r"flash-attention\.git@([0-9a-f]{40})"),
+        "fastvideo-kernel/README.md": _match(
+            "fastvideo-kernel/README.md",
+            r"flash-attention\.git@([0-9a-f]{40})"),
+    }
+
+    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

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 cutlass-4.5-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" }