[infra]: skip the FA4 cute overlay in the GB10 image

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.

.github/workflows/infra-build-image.yml

@@ -80,6 +80,9 @@ jobs:
             torch_cuda_arch_list:
               amd64: '9.0a'
               arm64: '9.0a'
+            install_fa4_cute:
+              amd64: '1'
+              arm64: '1'
           - version: '13.0.0'
             suffix: '-cuda13.0.0'
             torch_backend: 'cu130'
@@ -88,6 +91,11 @@ jobs:
             torch_cuda_arch_list:
               amd64: '9.0a'
               arm64: '12.1a'
+            # FA4 cute does not run on GB10 (sm_12x), and its cutlass-dsl pin
+            # breaks flashinfer there; GB10 FP4 attention uses fastvideo-kernel.
+            install_fa4_cute:
+              amd64: '1'
+              arm64: '0'
         architecture:
           - name: amd64
             runner: ubuntu-latest
@@ -107,6 +115,7 @@ jobs:
         CUDA_VERSION=${{ matrix.cuda.version }}
         UV_TORCH_BACKEND=${{ matrix.cuda.torch_backend }}
         TORCH_CUDA_ARCH_LIST=${{ matrix.cuda.torch_cuda_arch_list[matrix.architecture.name] }}
+        INSTALL_FA4_CUTE=${{ matrix.cuda.install_fa4_cute[matrix.architecture.name] }}
         CMAKE_BUILD_PARALLEL_LEVEL=${{ matrix.cuda.cmake_build_parallel_level }}
         FLASH_ATTN_WHEEL_TAG=${{ matrix.cuda.fa_tag }}
         FLASH_ATTN_WHEEL_RELEASE=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.17

docker/Dockerfile

@@ -23,7 +23,8 @@
 #   # 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 -t fastvideo-dev:spark .
+#       --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).
@@ -69,6 +70,11 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
 # 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
@@ -191,12 +197,18 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # 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')]" && \
-    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')"
+    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 . .
 

docs/contributing/developer_env/docker.md

@@ -37,7 +37,7 @@ native ARM64 host, build the CUDA 13 DGX Spark image with:
 
 ```bash
 docker build --platform linux/arm64 -f docker/Dockerfile \
-  --build-arg TORCH_CUDA_ARCH_LIST=12.1a \
+  --build-arg TORCH_CUDA_ARCH_LIST=12.1a --build-arg INSTALL_FA4_CUTE=0 \
   -t fastvideo-dev:spark .
 ```
 
@@ -45,6 +45,11 @@ This uses the prebuilt Linux ARM64 FlashAttention wheel. FastVideo's in-tree
 CUDA kernel is still compiled for `sm_121a` as part of the image build, including
 the FP4 attention kernels that `FASTVIDEO_H3_VSA_FP4=1` uses.
 
+`INSTALL_FA4_CUTE=0` (set for the published DGX Spark image) leaves
+out the FA4 `flash_attn.cute` overlay: FA4 supports SM 9/10/11 only, and its pinned
+`nvidia-cutlass-dsl` would break FlashInfer, which the NVFP4 paths need. FA2 is
+still installed.
+
 ## Using the container
 
 ```bash

fastvideo/tests/modal/test_kernel_build_cache.py

@@ -176,6 +176,27 @@ def test_ci_runner_image_targets_arm64_sm100_with_opencv_runtime() -> None:
     assert 'python -c "import cv2; print(\'OpenCV\', cv2.__version__)"' in dockerfile
 
 
+def test_gb10_image_builds_sm121a_without_fa4_cute() -> None:
+    """The CUDA 13 arm64 leg is the DGX Spark image: sm_121a keeps fastvideo-kernel's FP4
+    attention, and skipping the FA4 cute overlay keeps flashinfer importable."""
+    workflow = yaml.load(
+        (REPO_ROOT / ".github/workflows/infra-build-image.yml").read_text(encoding="utf-8"),
+        Loader=yaml.BaseLoader,
+    )
+    job = workflow["jobs"]["build-cuda-images"]
+    cuda = {entry["version"]: entry for entry in job["strategy"]["matrix"]["cuda"]}
+
+    assert cuda["13.0.0"]["torch_cuda_arch_list"]["arm64"] == "12.1a"
+    assert cuda["13.0.0"]["install_fa4_cute"] == {"amd64": "1", "arm64": "0"}
+    assert cuda["12.6.3"]["install_fa4_cute"] == {"amd64": "1", "arm64": "1"}
+    assert ("INSTALL_FA4_CUTE=${{ matrix.cuda.install_fa4_cute[matrix.architecture.name] }}"
+            in job["with"]["build_args"])
+
+    dockerfile = (REPO_ROOT / "docker/Dockerfile").read_text(encoding="utf-8")
+    assert "ARG INSTALL_FA4_CUTE=1" in dockerfile
+    assert 'if [ "${INSTALL_FA4_CUTE}" = "1" ]; then' in dockerfile
+
+
 def test_docker_image_bakes_modal_apt_and_rust_layer() -> None:
     """The dormant Modal rollback image layers apt packages + rustup on top of
     the release image (pr_test.py). Keep those packages baked into the release