From fba50356b1824ce4bc29ab7369dfba40c4176a3c Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Mon, 19 Oct 2026 08:13:04 +0000 Subject: [PATCH] [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 | 9 ++++++++ docker/Dockerfile | 18 +++++++++++++--- docs/contributing/developer_env/docker.md | 7 ++++++- .../tests/modal/test_kernel_build_cache.py | 21 +++++++++++++++++++ 4 files changed, 51 insertions(+), 4 deletions(-) diff --git a/.github/workflows/infra-build-image.yml b/.github/workflows/infra-build-image.yml index 31b548f88..d7940f7b6 100644 --- a/.github/workflows/infra-build-image.yml +++ b/.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 diff --git a/docker/Dockerfile b/docker/Dockerfile index 27d6c822e..2e85e5dbe 100644 --- a/docker/Dockerfile +++ b/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 . . diff --git a/docs/contributing/developer_env/docker.md b/docs/contributing/developer_env/docker.md index ee017d5f9..6208a3be7 100644 --- a/docs/contributing/developer_env/docker.md +++ b/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 diff --git a/fastvideo/tests/modal/test_kernel_build_cache.py b/fastvideo/tests/modal/test_kernel_build_cache.py index 077918564..92d91700c 100644 --- a/fastvideo/tests/modal/test_kernel_build_cache.py +++ b/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